repo_description: Go HTTP framework with high-performance and strong-extensibility for building micro-services.
repo_topics: go http microservices
code: /*
   * Copyright 2022 CloudWeGo Authors
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   *
   * The MIT License (MIT)
   *
   * Copyright (c) 2015-present Aliaksandr Valialkin, VertaMedia, Kirill Danshin, Erik Dubbelboer, FastHTTP Authors
   *
   * Permission is hereby granted, free of charge, to any person obtaining a copy
   * of this software and associated documentation files (the "Software"), to deal
   * in the Software without restriction, including without limitation the rights
   * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
   * copies of the Software, and to permit persons to whom the Software is
   * furnished to do so, subject to the following conditions:
   *
   * The above copyright notice and this permission notice shall be included in
   * all copies or substantial portions of the Software.
   *
   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
   * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
   * THE SOFTWARE.
   *
   * This file may have been modified by CloudWeGo authors. All CloudWeGo
   * Modifications are Copyright 2022 CloudWeGo Authors.
   */
  
  package protocol
  
  import (
  	"errors"
  	"io"
  	"net"
  	"sync"
  
  	"github.com/cloudwego/hertz/internal/bytesconv"
  	"github.com/cloudwego/hertz/internal/nocopy"
  	"github.com/cloudwego/hertz/pkg/common/bytebufferpool"
  	"github.com/cloudwego/hertz/pkg/common/compress"
  	"github.com/cloudwego/hertz/pkg/common/utils"
  	"github.com/cloudwego/hertz/pkg/network"
  )
  
  var (
  	responsePool sync.Pool
  	// NoResponseBody is an io.ReadCloser with no bytes. Read always returns EOF
  	// and Close always returns nil. It can be used in an ingoing client
  	// response to explicitly signal that a response has zero bytes.
  	NoResponseBody = noBody{}
  )
  
  // Response represents HTTP response.
  //
  // It is forbidden copying Response instances. Create new instances
  // and use CopyTo instead.
  //
  // Response instance MUST NOT be used from concurrently running goroutines.
  type Response struct {
  	noCopy nocopy.NoCopy //lint:ignore U1000 until noCopy is used
  
  	// Response header
  	//
  	// Copying Header by value is forbidden. Use pointer to Header instead.
  	Header ResponseHeader
  
  	// Flush headers as soon as possible without waiting for first body bytes.
  	// Relevant for bodyStream only.
  	ImmediateHeaderFlush bool
  
  	bodyStream      io.Reader
  	w               responseBodyWriter
  	body            *bytebufferpool.ByteBuffer
  	bodyRaw         []byte
  	maxKeepBodySize int
  
  	// Response.Read() skips reading body if set to true.
  	// Use it for reading HEAD responses.
  	//
  	// Response.Write() skips writing body if set to true.
  	// Use it for writing HEAD responses.
  	SkipBody bool
  
  	// Remote TCPAddr from concurrently net.Conn
  	raddr net.Addr
  	// Local TCPAddr from concurrently net.Conn
  	laddr net.Addr
  
  	// If set a hijackWriter, hertz will skip the default header/body writer process.
  	hijackWriter network.ExtWriter
  }
  
  func (resp *Response) GetHijackWriter() network.ExtWriter {
  	return resp.hijackWriter
  }
  
  func (resp *Response) HijackWriter(writer network.ExtWriter) {
  	resp.hijackWriter = writer
  }
  
  type responseBodyWriter struct {
  	r *Response
  }
  
  func (w *responseBodyWriter) Write(p []byte) (int, error) {
  	w.r.AppendBody(p)
  	return len(p), nil
  }
  
  func (resp *Response) MustSkipBody() bool {
  	return resp.SkipBody || resp.Header.MustSkipContentLength()
  }
  
  // BodyGunzip returns un-gzipped body data.
  //
  // This method may be used if the response header contains
  // 'Content-Encoding: gzip' for reading un-gzipped body.
  // Use Body for reading gzipped response body.
  func (resp *Response) BodyGunzip() ([]byte, error) {
  	return gunzipData(resp.Body())
  }
  
  // SetConnectionClose sets 'Connection: close' header.
  func (resp *Response) SetConnectionClose() {
  	resp.Header.SetConnectionClose(true)
  }
  
  // SetBodyString sets response body.
  func (resp *Response) SetBodyString(body string) {
  	resp.CloseBodyStream()            //nolint:errcheck
  	resp.BodyBuffer().SetString(body) //nolint:errcheck
  }
  
  func (resp *Response) ConstructBodyStream(body *bytebufferpool.ByteBuffer, bodyStream io.Reader) {
  	resp.body = body
  	resp.bodyStream = bodyStream
  }
  
  // BodyWriter returns writer for populating response body.
  //
  // If used inside RequestHandler, the returned writer must not be used
  // after returning from RequestHandler. Use RequestContext.Write
  // or SetBodyStreamWriter in this case.
  func (resp *Response) BodyWriter() io.Writer {
  	resp.w.r = resp
  	return &resp.w
  }
  
  // SetStatusCode sets response status code.
  func (resp *Response) SetStatusCode(statusCode int) {
  	resp.Header.SetStatusCode(statusCode)
  }
  
  func (resp *Response) SetMaxKeepBodySize(n int) {
  	resp.maxKeepBodySize = n
  }
  
  func (resp *Response) BodyBytes() []byte {
  	if resp.bodyRaw != nil {
  		return resp.bodyRaw
  	}
  	if resp.body == nil {
  		return nil
  	}
  	return resp.body.B
  }
  
  func (resp *Response) HasBodyBytes() bool {
  	return len(resp.BodyBytes()) != 0
  }
  
  func (resp *Response) CopyToSkipBody(dst *Response) {
  	dst.Reset()
  	resp.Header.CopyTo(&dst.Header)
  	dst.SkipBody = resp.SkipBody
  	dst.raddr = resp.raddr
  	dst.laddr = resp.laddr
  }
  
  // IsBodyStream returns true if body is set via SetBodyStream*
  func (resp *Response) IsBodyStream() bool {
  	return resp.bodyStream != nil
  }
  
  // SetBodyStream sets response body stream and, optionally body size.
  //
  // If bodySize is >= 0, then the bodyStream must provide exactly bodySize bytes
  // before returning io.EOF.
  //
  // If bodySize < 0, then bodyStream is read until io.EOF.
  //
  // bodyStream.Close() is called after finishing reading all body data
  // if it implements io.Closer.
  //
  // See also SetBodyStreamWriter.
  func (resp *Response) SetBodyStream(bodyStream io.Reader, bodySize int) {
  	resp.ResetBody()
  	resp.bodyStream = bodyStream
  	resp.Header.SetContentLength(bodySize)
  }
  
  // SetBodyStreamNoReset is almost the same as SetBodyStream,
  // but it doesn't reset the bodyStream before.
  func (resp *Response) SetBodyStreamNoReset(bodyStream io.Reader, bodySize int) {
  	resp.bodyStream = bodyStream
  	resp.Header.SetContentLength(bodySize)
  }
  
  // BodyE returns response body.
  func (resp *Response) BodyE() ([]byte, error) {
  	if resp.bodyStream != nil {
  		bodyBuf := resp.BodyBuffer()
  		bodyBuf.Reset()
  		zw := network.NewWriter(bodyBuf)
  		_, err := utils.CopyZeroAlloc(zw, resp.bodyStream)
  		resp.CloseBodyStream() //nolint:errcheck
  		if err != nil {
  			return nil, err
  		}
  	}
  	return resp.BodyBytes(), nil
  }
  
  // Body returns response body.
  // if get body failed, returns nil.
  func (resp *Response) Body() []byte {
  	body, _ := resp.BodyE()
  	return body
  }
  
  // BodyWriteTo writes response body to w.
  func (resp *Response) BodyWriteTo(w io.Writer) error {
  	zw := network.NewWriter(w)
  	if resp.bodyStream != nil {
  		_, err := utils.CopyZeroAlloc(zw, resp.bodyStream)
  		resp.CloseBodyStream() //nolint:errcheck
  		return err
  	}
  
  	body := resp.BodyBytes()
  	zw.WriteBinary(body) //nolint:errcheck
  	return zw.Flush()
  }
  
  // CopyTo copies resp contents to dst except of body stream.
  func (resp *Response) CopyTo(dst *Response) {
  	resp.CopyToSkipBody(dst)
  	if resp.bodyRaw != nil {
  		dst.bodyRaw = append(dst.bodyRaw[:0], resp.bodyRaw...)
  		if dst.body != nil {
  			dst.body.Reset()
  		}
  	} else if resp.body != nil {
  		dst.BodyBuffer().Set(resp.body.B)
  	} else if dst.body != nil {
  		dst.body.Reset()
  	}
  }
  
  func SwapResponseBody(a, b *Response) {
  	a.body, b.body = b.body, a.body
  	a.bodyRaw, b.bodyRaw = b.bodyRaw, a.bodyRaw
  	a.bodyStream, b.bodyStream = b.bodyStream, a.bodyStream
  }
  
  // Reset clears response contents.
  func (resp *Response) Reset() {
  	resp.Header.Reset()
  	resp.resetSkipHeader()
  	resp.SkipBody = false
  	resp.raddr = nil
  	resp.laddr = nil
  	resp.ImmediateHeaderFlush = false
  	resp.hijackWriter = nil
  }
  
  func (resp *Response) resetSkipHeader() {
  	resp.ResetBody()
  }
  
  // ResetBody resets response body.
  func (resp *Response) ResetBody() {
  	resp.bodyRaw = nil
  	resp.CloseBodyStream() //nolint:errcheck
  	if resp.body != nil {
  		if resp.body.Cap() <= resp.maxKeepBodySize {
  			resp.body.Reset()
  			return
  		}
  		responseBodyPool.Put(resp.body)
  		resp.body = nil
  	}
  }
  
  // SetBodyRaw sets response body, but without copying it.
  //
  // From this point onward the body argument must not be changed.
  func (resp *Response) SetBodyRaw(body []byte) {
  	resp.ResetBody()
  	resp.bodyRaw = body
  }
  
  // StatusCode returns response status code.
  func (resp *Response) StatusCode() int {
  	return resp.Header.StatusCode()
  }
  
  // SetBody sets response body.
  //
  // It is safe re-using body argument after the function returns.
  func (resp *Response) SetBody(body []byte) {
  	resp.CloseBodyStream() //nolint:errcheck
  	if resp.GetHijackWriter() == nil {
  		resp.BodyBuffer().Set(body) //nolint:errcheck
  		return
  	}
  
  	// If the hijack writer support .SetBody() api, then use it.
  	if setter, ok := resp.GetHijackWriter().(interface {
  		SetBody(b []byte)
  	}); ok {
  		setter.SetBody(body)
  		return
  	}
  
  	// Otherwise, call .Write() api instead.
  	resp.GetHijackWriter().Write(body) //nolint:errcheck
  }
  
  func (resp *Response) BodyStream() io.Reader {
  	if resp.bodyStream == nil {
  		return NoResponseBody
  	}
  	return resp.bodyStream
  }
  
  // Hijack returns the underlying network.Conn if available.
  //
  // It's only available when StatusCode() == 101 and "Connection: Upgrade",
  // coz Hertz will NOT reuse connection in this case,
  // then make it optional for users to implement their own protocols.
  //
  // The most common scenario is used with github.com/hertz-contrib/websocket
  func (resp *Response) Hijack() (network.Conn, error) {
  	if resp.bodyStream != nil {
  		h, ok := resp.bodyStream.(interface {
  			Hijack() (network.Conn, error)
  		})
  		if ok {
  			return h.Hijack()
  		}
  	}
  	return nil, errors.New("not available")
  }
  
  // AppendBody appends p to response body.
  //
  // It is safe re-using p after the function returns.
  func (resp *Response) AppendBody(p []byte) {
  	resp.CloseBodyStream() //nolint:errcheck
  	if resp.hijackWriter != nil {
  		resp.hijackWriter.Write(p) //nolint:errcheck
  		return
  	}
  	resp.BodyBuffer().Write(p) //nolint:errcheck
  }
  
  // AppendBodyString appends s to response body.
  func (resp *Response) AppendBodyString(s string) {
  	resp.CloseBodyStream() //nolint:errcheck
  	if resp.hijackWriter != nil {
  		resp.hijackWriter.Write(bytesconv.S2b(s)) //nolint:errcheck
  		return
  	}
  	resp.BodyBuffer().WriteString(s) //nolint:errcheck
  }
  
  // ConnectionClose returns true if 'Connection: close' header is set.
  func (resp *Response) ConnectionClose() bool {
  	return resp.Header.ConnectionClose()
  }
  
  // CloseBodyStream tries call Close() of underlying body stream.
  //
  // NOTE:
  // * MUST NOT call CloseBodyStream() and BodyStream().Read() concurrently to avoid race issue.
  func (resp *Response) CloseBodyStream() error {
  	if resp.bodyStream == nil {
  		return nil
  	}
  	var err error
  	if bsc, ok := resp.bodyStream.(io.Closer); ok {
  		err = bsc.Close()
  	}
  	resp.bodyStream = nil
  	return err
  }
  
  func (resp *Response) BodyBuffer() *bytebufferpool.ByteBuffer {
  	if resp.body == nil {
  		resp.body = responseBodyPool.Get()
  	}
  	resp.bodyRaw = nil
  	return resp.body
  }
  
  func gunzipData(p []byte) ([]byte, error) {
  	var bb bytebufferpool.ByteBuffer
  	_, err := compress.WriteGunzip(&bb, p)
  	if err != nil {
  		return nil, err
  	}
  	return bb.B, nil
  }
  
  // RemoteAddr returns the remote network address. The Addr returned is shared
  // by all invocations of RemoteAddr, so do not modify it.
  func (resp *Response) RemoteAddr() net.Addr {
  	return resp.raddr
  }
  
  // LocalAddr returns the local network address. The Addr returned is shared
  // by all invocations of LocalAddr, so do not modify it.
  func (resp *Response) LocalAddr() net.Addr {
  	return resp.laddr
  }
  
  func (resp *Response) ParseNetAddr(conn network.Conn) {
  	resp.raddr = conn.RemoteAddr()
  	resp.laddr = conn.LocalAddr()
  }
  
  // AcquireResponse returns an empty Response instance from response pool.
  //
  // The returned Response instance may be passed to ReleaseResponse when it is
  // no longer needed. This allows Response recycling, reduces GC pressure
  // and usually improves performance.
  func AcquireResponse() *Response {
  	v := responsePool.Get()
  	if v == nil {
  		return &Response{}
  	}
  	return v.(*Response)
  }
  
  // ReleaseResponse return resp acquired via AcquireResponse to response pool.
  //
  // It is forbidden accessing resp and/or its members after returning
  // it to response pool.
  func ReleaseResponse(resp *Response) {
  	resp.Reset()
  	responsePool.Put(resp)
  }
  