code.vegaprotocol.io/vega@v0.79.0/datanode/gateway/rest/gziphandler.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package rest
    17  
    18  import (
    19  	"compress/gzip"
    20  	"net/http"
    21  	"strings"
    22  	"sync"
    23  
    24  	"code.vegaprotocol.io/vega/logging"
    25  )
    26  
    27  type gzipResponseWriter struct {
    28  	http.ResponseWriter
    29  
    30  	w             *gzip.Writer
    31  	statusCode    int
    32  	headerWritten bool
    33  }
    34  
    35  var pool = sync.Pool{
    36  	New: func() interface{} {
    37  		w, _ := gzip.NewWriterLevel(nil, gzip.BestSpeed)
    38  		return &gzipResponseWriter{
    39  			w: w,
    40  		}
    41  	},
    42  }
    43  
    44  func (gzr *gzipResponseWriter) WriteHeader(statusCode int) {
    45  	gzr.statusCode = statusCode
    46  	gzr.headerWritten = true
    47  
    48  	if gzr.statusCode != http.StatusNotModified && gzr.statusCode != http.StatusNoContent {
    49  		gzr.ResponseWriter.Header().Del("Content-Length")
    50  		gzr.ResponseWriter.Header().Set("Content-Encoding", "gzip")
    51  	}
    52  
    53  	gzr.ResponseWriter.WriteHeader(statusCode)
    54  }
    55  
    56  func (gzr *gzipResponseWriter) Write(b []byte) (int, error) {
    57  	if _, ok := gzr.Header()["Content-Type"]; !ok {
    58  		// If no content type, apply sniffing algorithm to un-gzipped body.
    59  		gzr.ResponseWriter.Header().Set("Content-Type", http.DetectContentType(b))
    60  	}
    61  
    62  	if !gzr.headerWritten {
    63  		// This is exactly what Go would also do if it hasn't been written yet.
    64  		gzr.WriteHeader(http.StatusOK)
    65  	}
    66  
    67  	return gzr.w.Write(b)
    68  }
    69  
    70  func (gzr *gzipResponseWriter) Flush() {
    71  	if gzr.w != nil {
    72  		gzr.w.Flush()
    73  	}
    74  
    75  	if fw, ok := gzr.ResponseWriter.(http.Flusher); ok {
    76  		fw.Flush()
    77  	}
    78  }
    79  
    80  func NewGzipHandler(logger logging.Logger, fn http.HandlerFunc) http.HandlerFunc {
    81  	return func(w http.ResponseWriter, r *http.Request) {
    82  		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
    83  			fn(w, r)
    84  			return
    85  		}
    86  
    87  		gzr := pool.Get().(*gzipResponseWriter)
    88  		gzr.statusCode = 0
    89  		gzr.headerWritten = false
    90  		gzr.ResponseWriter = w
    91  		gzr.w.Reset(w)
    92  
    93  		defer func() {
    94  			// gzr.w.Close will write a footer even if no data has been written.
    95  			// StatusNotModified and StatusNoContent expect an empty body so don't close it.
    96  			if gzr.statusCode != http.StatusNotModified && gzr.statusCode != http.StatusNoContent {
    97  				if err := gzr.w.Close(); err != nil {
    98  					logger.Error("Failed to Gzip output from REST proxy", logging.Error(err))
    99  				}
   100  			}
   101  			pool.Put(gzr)
   102  		}()
   103  
   104  		fn(gzr, r)
   105  	}
   106  }