github.com/zhongdalu/gf@v1.0.0/g/net/ghttp/ghttp_response_writer.go (about)

     1  // Copyright 2017 gf Author(https://github.com/zhongdalu/gf). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/zhongdalu/gf.
     6  //
     7  
     8  package ghttp
     9  
    10  import (
    11  	"bytes"
    12  	"net/http"
    13  )
    14  
    15  // 自定义的ResponseWriter，用于写入流的控制
    16  type ResponseWriter struct {
    17  	http.ResponseWriter
    18  	Status int           // http status
    19  	buffer *bytes.Buffer // 缓冲区内容
    20  }
    21  
    22  // 覆盖父级的WriteHeader方法
    23  func (w *ResponseWriter) Write(data []byte) (int, error) {
    24  	w.buffer.Write(data)
    25  	return len(data), nil
    26  }
    27  
    28  // 覆盖父级的WriteHeader方法, 这里只会记录Status做缓冲处理, 并不会立即输出到HEADER。
    29  func (w *ResponseWriter) WriteHeader(status int) {
    30  	w.Status = status
    31  }
    32  
    33  // 输出buffer数据到客户端.
    34  func (w *ResponseWriter) OutputBuffer() {
    35  	if w.Status != 0 {
    36  		w.ResponseWriter.WriteHeader(w.Status)
    37  	}
    38  	if w.buffer.Len() > 0 {
    39  		w.ResponseWriter.Write(w.buffer.Bytes())
    40  		w.buffer.Reset()
    41  	}
    42  }