まえがき

Goのechoフレームワークでトレースログとエラーコードをロギングしようとしたら、エラーコードを決定するロジックが複数あり、レスポンスから取得するのが一番影響が少ないことに気づいてしまった困った人向け

エラー時のみハンドリングしたかったり、トレースログを含める必要があったので、 middleware/body-dump をカスタマイズする方向でミドルウェアを書いた。

トリッキーな方法だと思うので、単純に Response Body を取得したいだけであれば middleware/body-dump を使うことを推奨

仕様とかそこらへん

Response Body を取得したいけど Response は stream に書き込まれてるっぽくって後から取得しようとしてもできなかったので、Context.Response().Writer をトラップ付き Writer への参照に書き換えて stream と memory の両方へ書き込まれるようにした。

echo/middleware/body_dump.go - GitHub をベースに~~車輪の再発明~~カスタマイズ。

必要なときにハンドラを設定して Response が書き込まれたら設定しておいた handler が呼ばれる仕様とした。

コアとなる部分の実装

package handler

import (
	"io"
	"fmt"
	"bytes"
	"bufio"
	"net"
	"net/http"

	"github.com/labstack/echo"
)

// ここから
type(
	bodyDumpResponseWriter struct {
		io.Writer
		http.ResponseWriter
	}
)

func (w *bodyDumpResponseWriter) WriteHeader(code int) {
	w.ResponseWriter.WriteHeader(code)
}

func (w *bodyDumpResponseWriter) Write(b []byte) (int, error) {
	return w.Writer.Write(b)
}

func (w *bodyDumpResponseWriter) Flush() {
	w.ResponseWriter.(http.Flusher).Flush()
}

func (w *bodyDumpResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return w.ResponseWriter.(http.Hijacker).Hijack()
}
// ここまで middleware/body-dump の内容まま

// 以降が記事の処理の実体
func ResponseHandler(c echo.Context, err error) {
    // Body を memory へ dump する為の Buffer
	resBody := new(bytes.Buffer)

	// http.ResponseWriter と 用意した Buffer の両方へ書き込むために io.MultiWriter を使用
	mw := io.MultiWriter(c.Response().Writer, resBody)

	// Body の dump 機能付き Writer の作成
	writer := &bodyDumpResponseWriter{Writer: mw, ResponseWriter: c.Response().Writer}

	// Context の Writer を先程作成した Writer に置換
	c.Response().Writer = writer

	// Response が書き込まれた際に呼ばれる関数を set
	c.Response().After(responseHandler(c, err, resBody))
}

func responseHandler(c echo.Context, err error, resBody *bytes.Buffer) func() {
	return func() {
		// ここで resBody のパースや error の分解や Context の内容を取得する処理を定義

		fmt.Printf("%s: %s", err, resBody.String())
	}
}

使い方

package main

import (
	"net/http"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"

	"./handler"
)

func main() {
	e := echo.New()

	e.Use(middleware.Logger())
	e.Use(middleware.Recover())

	e.GET("/", hello)

	e.Logger.Fatal(e.Start(":1323"))
}

func hello(c echo.Context) error {
	// Response が書き込まれる直前に設定
	handler.ResponseHandler(c, nil)

	// Context が Response を書き込み終えると responseHandler() が call される
	return c.String(http.StatusOK, "Hello, World!")
}

エラーハンドラの中でエラーレスポンスを生成してる箇所とかあればその直前にセットすれば良いと思う。

// JSONを返すコードの場合はこんな感じ
handler.ResponseHandler(c, nil)
c.JSON(http.StatusInternalServerError, ...)

注意

echo - GoDoc Response.After

After registers a function which is called just after the response is written. If the Content-Length is unknown, none of the after function is executed.

今回のケースでは問題とならないが、 Response.After の注意書きに「Contents-Length が不明な場合は呼ばれないよ。」と書いてある。ファイルのアップストリームとかでは使えないようだ

参考にした資料