golangの日記

Go言語を中心にプログラミングについてのブログ

Go言語(golang)で httpリクエストのレスポンスボディの大きさにリミットを設ける

golang.png


HTTP Request で返ってくる http.Response の Body の大きさに制限を設けてコンテンツが制限を超えていればリクエストを中止する方法です。
レスポンスヘッダーの Content-Length または、Go言語の場合は、resp.ContentLength をもとに制限すればよいのかもしれませんが、
そもそもレスポンスヘッダーに Content-Length がないため -1 になっていたりして、よくわからないので、他の方法でリミットを設けたい。





レスポンスボディを ioutil.ReadAll で読み出す前に io.LimitedReader を元に、
それに少し変更を加えたものを挟んで、その Read 関数内でコンテンツ長を数えています。
データの大きさが limit を超えてた場合は、その時点で errors.New("Exceeded") のエラーになります。


// 制限を設けるために ioutil.ReadAll の前に Read するためのもの
type LimitedReadCloser struct {
    RC io.ReadCloser
    N  int64
}

func LimitReadCloser(r io.ReadCloser, n int64) io.ReadCloser {
    return &LimitedReadCloser{r, n}
}

func (l *LimitedReadCloser) Close() error {
    return l.RC.Close()
}

func (l *LimitedReadCloser) Read(p []byte) (n int, err error) {
    if l.N <= 0 {
        return 0, errors.New("Exceeded")
    }
    if int64(len(p)) > l.N {
        p = p[0:l.N]
    }
    n, err = l.RC.Read(p)
    l.N -= int64(n)
    return
}

func main() {
    resp, err := http.Get("https://example.com")
    if err != nil {
        log.Fatal(err)
    }
    defer resp.Body.Close()

    const limit = 5 << 20 // 約5MB

    // 第二引数に制限するバイト数を指定
    resp.Body = LimitReadCloser(resp.Body, limit)

    data, err := ioutil.ReadAll(resp.Body)
    if err != nil {
        log.Fatal(err)
    }

    fmt.Println(len(data))
}