github.com/cloudwego/hertz@v0.9.3/pkg/common/adaptor/response.go (about)

     1  /*
     2   * Copyright 2022 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package adaptor
    18  
    19  import (
    20  	"net/http"
    21  
    22  	"github.com/cloudwego/hertz/pkg/protocol"
    23  	"github.com/cloudwego/hertz/pkg/protocol/consts"
    24  )
    25  
    26  type compatResponse struct {
    27  	h           *protocol.Response
    28  	header      http.Header
    29  	writeHeader bool
    30  }
    31  
    32  func (c *compatResponse) Header() http.Header {
    33  	if c.header != nil {
    34  		return c.header
    35  	}
    36  	c.header = make(map[string][]string)
    37  	return c.header
    38  }
    39  
    40  func (c *compatResponse) Write(b []byte) (int, error) {
    41  	if !c.writeHeader {
    42  		c.WriteHeader(consts.StatusOK)
    43  	}
    44  
    45  	return c.h.BodyWriter().Write(b)
    46  }
    47  
    48  func (c *compatResponse) WriteHeader(statusCode int) {
    49  	if !c.writeHeader {
    50  		for k, v := range c.header {
    51  			for _, vv := range v {
    52  				if k == consts.HeaderContentLength {
    53  					continue
    54  				}
    55  				if k == consts.HeaderSetCookie {
    56  					cookie := protocol.AcquireCookie()
    57  					_ = cookie.Parse(vv)
    58  					c.h.Header.SetCookie(cookie)
    59  					continue
    60  				}
    61  				c.h.Header.Add(k, vv)
    62  			}
    63  		}
    64  		c.writeHeader = true
    65  	}
    66  
    67  	c.h.Header.SetStatusCode(statusCode)
    68  }
    69  
    70  // GetCompatResponseWriter only support basic function of ResponseWriter, not for all.
    71  func GetCompatResponseWriter(resp *protocol.Response) http.ResponseWriter {
    72  	c := &compatResponse{
    73  		h: resp,
    74  	}
    75  	c.h.Header.SetNoDefaultContentType(true)
    76  
    77  	h := make(map[string][]string)
    78  	tmpKey := make([][]byte, 0, c.h.Header.Len())
    79  	c.h.Header.VisitAll(func(k, v []byte) {
    80  		h[string(k)] = append(h[string(k)], string(v))
    81  		tmpKey = append(tmpKey, k)
    82  	})
    83  
    84  	for _, k := range tmpKey {
    85  		c.h.Header.DelBytes(k)
    86  	}
    87  
    88  	c.header = h
    89  	return c
    90  }