はじめに

https://mattn.kaoriya.net/software/lang/go/20211224005655.htm
こちらの記事を読んだので中身を追ってみました

概要

go.1.18でhttp.MaxBytesHandlerが入った
https://pkg.go.dev/net/http@go1.18beta1#MaxBytesHandler

MaxBytesHandler returns a Handler that runs h with its ResponseWriter and Request.Body wrapped by a MaxBytesReader.

この機能はもともとMaxBytesReaderという機能が存在していて
それをラップするハンドラーとして追加されたようです

使用方法

middlewareとして実装するようになっています
- こういう自前のラップが嫌な場合はrouterがそこそこあるので
- そちらを使用したほうがいいと思います
  - https://pkg.go.dev/search?q=router
  - ここ以外にもweb framworkに同梱されているものもあります

package main

import (
	"log"
	"net/http"
)

func middleware(next http.Handler) http.Handler {
	return http.MaxBytesHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		next.ServeHTTP(w, r)
	}), 4096)
}

func main() {
	http.Handle("/", middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("OK"))
	})))

	log.Fatal(http.ListenAndServe(":8080", nil))
}

MaxBytesHandlerの処理

内部的にはBodyをMaxBytesReaderで読み込んでいるだけです

func MaxBytesHandler(h Handler, n int64) Handler {  
   return HandlerFunc(func(w ResponseWriter, r *Request) {  
      r2 := *r  
      r2.Body = MaxBytesReader(w, r.Body, n)  
      h.ServeHTTP(w, &r2)  
   })  
}

MaxBytesReader.Read

Readerのinterfaceを実装しているので
実質上の実装はReadになります

func (l *maxBytesReader) Read(p []byte) (n int, err error) {
	if l.err != nil {
		return 0, l.err
	}
	if len(p) == 0 {
		return 0, nil
	}
	// If they asked for a 32KB byte read but only 5 bytes are
	// remaining, no need to read 32KB. 6 bytes will answer the
	// question of the whether we hit the limit or go past it.
	if int64(len(p)) > l.n+1 {
		p = p[:l.n+1]
	}
	n, err = l.r.Read(p)

	if int64(n) <= l.n {
		l.n -= int64(n)
		l.err = err
		return n, err
	}

	n = int(l.n)
	l.n = 0

	// The server code and client code both use
	// maxBytesReader. This "requestTooLarge" check is
	// only used by the server code. To prevent binaries
	// which only using the HTTP Client code (such as
	// cmd/go) from also linking in the HTTP server, don't
	// use a static type assertion to the server
	// "*response" type. Check this interface instead:
	type requestTooLarger interface {
		requestTooLarge()
	}
	if res, ok := l.w.(requestTooLarger); ok {
		res.requestTooLarge()
	}
	l.err = errors.New("http: request body too large")
	return n, l.err
}

limitより大きい場合はサイズ内に詰め替え

無駄な読み込みを減らすためだとはおもいますが
制限値より大きい場合はsliceをサイズ内に修正してます

// If they asked for a 32KB byte read but only 5 bytes are
// remaining, no need to read 32KB. 6 bytes will answer the
// question of the whether we hit the limit or go past it.
if int64(len(p)) > l.n+1 {
	p = p[:l.n+1]
}

これ簡単に例を示すと
制限： 3
len：5
[0,1,2,3,4] -> [:n+1] -> [:4] -> [0,1,2,3]

制限より読み込みが少ない場合はそのまま返す

ここが一番あたま抱えたんですが
データが制限よりすくない場合
全部エラーになるとおもいこんでいたんですよね
実際にリクエスト飛ばすとOKが返ってくる
見た目のまんま、制限より少ない場合はそのままreturnしているだけでした

n, err = l.r.Read(p)

if int64(n) <= l.n {
	l.n -= int64(n)
	l.err = err
	return n, err
}

制限より大きかった場合

ここから先は制限より大きかった場合の話ですが
ここでwが何かを確認するとResponseWriterが入っています

type maxBytesReader struct {
	w   ResponseWriter
	r   io.ReadCloser // underlying reader
	n   int64         // max bytes remaining
	err error         // sticky error
}

受け取ったResponseWirterがrequestTooLargerを実装していた場合
このメソッドが走ります

type requestTooLarger interface {
	requestTooLarge()
}
if res, ok := l.w.(requestTooLarger); ok {
	res.requestTooLarge()
}

このメソッドは何をしているかといいますと
responseのcloseAfterReplyとrequestBodyLimitHitをtrueに設定して
HeaderにConnection:closeをセットしています

// requestTooLarge is called by maxBytesReader when too much input has
// been read from the client.
func (w *response) requestTooLarge() {
	w.closeAfterReply = true
	w.requestBodyLimitHit = true
	if !w.wroteHeader {
		w.Header().Set("Connection", "close")
	}
}

それぞれの定義は以下のようになります
keepaliveとかを考慮されたもののようですね

// close connection after this reply.  set on request and
// updated after response from handler if there's a
// "Connection: keep-alive" response header and a
// Content-Length.
closeAfterReply bool

// requestBodyLimitHit is set by requestTooLarge when
// maxBytesReader hits its max size. It is checked in
// WriteHeader, to make sure we don't consume the
// remaining request body to try to advance to the next HTTP
// request. Instead, when this is set, we stop reading
// subsequent requests on this connection and stop reading
// input from it.
requestBodyLimitHit bool

エラーを返す

最終的にhttp: request body too largeが設定されて
returnされます

l.err = errors.New("http: request body too large")

Discussion