trpc.group/trpc-go/trpc-go@v1.0.3/restful/compressor.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package restful
    15  
    16  import (
    17  	"io"
    18  	"net/http"
    19  )
    20  
    21  // Compressor is the interface for http body compression/decompression.
    22  type Compressor interface {
    23  	// Compress compresses http body.
    24  	Compress(w io.Writer) (io.WriteCloser, error)
    25  	// Decompress decompresses http body.
    26  	Decompress(r io.Reader) (io.Reader, error)
    27  	// Name returns name of the Compressor.
    28  	Name() string
    29  	// ContentEncoding returns the encoding indicated by Content-Encoding response header.
    30  	ContentEncoding() string
    31  }
    32  
    33  // compressor related http headers
    34  var (
    35  	headerAcceptEncoding  = http.CanonicalHeaderKey("Accept-Encoding")
    36  	headerContentEncoding = http.CanonicalHeaderKey("Content-Encoding")
    37  )
    38  
    39  var compressors = make(map[string]Compressor)
    40  
    41  // RegisterCompressor registers a Compressor.
    42  // This function is not thread-safe, it should only be called in init() function.
    43  func RegisterCompressor(c Compressor) {
    44  	if c == nil || c.Name() == "" {
    45  		panic("tried to register nil or anonymous compressor")
    46  	}
    47  	compressors[c.Name()] = c
    48  }
    49  
    50  // GetCompressor returns a Compressor by name.
    51  func GetCompressor(name string) Compressor {
    52  	return compressors[name]
    53  }
    54  
    55  // compressorForTranscoding returns inbound/outbound Compressors for transcoding.
    56  func compressorForTranscoding(contentEncodings []string, acceptEncodings []string) (Compressor, Compressor) {
    57  	var reqCompressor, respCompressor Compressor // both could be nil
    58  
    59  	for _, contentEncoding := range contentEncodings {
    60  		if c, ok := compressors[contentEncoding]; ok {
    61  			reqCompressor = c
    62  			break
    63  		}
    64  	}
    65  
    66  	for _, acceptEncoding := range acceptEncodings {
    67  		if c, ok := compressors[acceptEncoding]; ok {
    68  			respCompressor = c
    69  			break
    70  		}
    71  	}
    72  
    73  	return reqCompressor, respCompressor
    74  }