github.com/alibaba/ilogtail/pkg@v0.0.0-20250526110833-c53b480d046c/protocol/decoder/common/comon.go (about)

     1  // Copyright 2021 iLogtail Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package common
    16  
    17  import (
    18  	"bytes"
    19  	"compress/gzip"
    20  	"errors"
    21  	"fmt"
    22  	"io"
    23  	"net/http"
    24  	"strconv"
    25  	"sync"
    26  
    27  	"github.com/golang/snappy"
    28  	"github.com/pierrec/lz4"
    29  )
    30  
    31  const (
    32  	ProtocolSLS          = "sls"
    33  	ProtocolPrometheus   = "prometheus"
    34  	ProtocolInflux       = "influx"
    35  	ProtocolInfluxdb     = "influxdb"
    36  	ProtocolStatsd       = "statsd"
    37  	ProtocolOTLPLogV1    = "otlp_logv1"
    38  	ProtocolOTLPMetricV1 = "otlp_metricv1"
    39  	ProtocolOTLPTraceV1  = "otlp_tracev1"
    40  	ProtocolRaw          = "raw"
    41  	ProtocolPyroscope    = "pyroscope"
    42  )
    43  
    44  var bufPool = sync.Pool{
    45  	New: func() interface{} {
    46  		buf := bytes.NewBuffer(make([]byte, 0, 32*1024))
    47  		return buf
    48  	},
    49  }
    50  
    51  func GetPooledBuf() *bytes.Buffer {
    52  	buf := bufPool.Get().(*bytes.Buffer)
    53  	return buf
    54  }
    55  
    56  func PutPooledBuf(buf *bytes.Buffer) {
    57  	buf.Reset()
    58  	bufPool.Put(buf)
    59  }
    60  
    61  func CollectBody(res http.ResponseWriter, req *http.Request, maxBodySize int64) ([]byte, int, error) {
    62  	body := req.Body
    63  
    64  	// Handle gzip request bodies
    65  	if req.Header.Get("Content-Encoding") == "gzip" {
    66  		var err error
    67  		body, err = gzip.NewReader(req.Body)
    68  		if err != nil {
    69  			return nil, http.StatusBadRequest, err
    70  		}
    71  		defer body.Close()
    72  	}
    73  
    74  	body = http.MaxBytesReader(res, body, maxBodySize)
    75  
    76  	if req.Header.Get("Content-Encoding") == "snappy" {
    77  		// for snappy encoding, use pooled buf to read compressed request body
    78  		buf := GetPooledBuf()
    79  		defer PutPooledBuf(buf)
    80  		_, err := io.Copy(buf, body) // nolint
    81  		if err != nil {
    82  			return nil, http.StatusBadRequest, err
    83  		}
    84  		data, err := snappy.Decode(nil, buf.Bytes())
    85  		if err != nil {
    86  			return nil, http.StatusBadRequest, err
    87  		}
    88  		return data, http.StatusOK, nil
    89  	}
    90  
    91  	bytes, err := io.ReadAll(body)
    92  	if err != nil {
    93  		return nil, http.StatusRequestEntityTooLarge, err
    94  	}
    95  
    96  	if req.Header.Get("x-log-compresstype") == "lz4" {
    97  		rawBodySize, err := strconv.Atoi(req.Header.Get("x-log-bodyrawsize"))
    98  		if err != nil || rawBodySize <= 0 {
    99  			return nil, http.StatusBadRequest, errors.New("invalid x-log-compresstype header " + req.Header.Get("x-log-bodyrawsize"))
   100  		}
   101  		data := make([]byte, rawBodySize)
   102  		if readSize, err := lz4.UncompressBlock(bytes, data); readSize != rawBodySize || (err != nil && err != io.EOF) {
   103  			return nil, http.StatusBadRequest, fmt.Errorf("uncompress lz4 error, expect : %d, real : %d, error : %v ", readSize, rawBodySize, err)
   104  		}
   105  		return data, http.StatusOK, nil
   106  	}
   107  
   108  	return bytes, http.StatusOK, nil
   109  }
   110  
   111  func CollectRawBody(res http.ResponseWriter, req *http.Request, maxBodySize int64) ([]byte, int, error) {
   112  	body := req.Body
   113  	body = http.MaxBytesReader(res, body, maxBodySize)
   114  	bytes, err := io.ReadAll(body)
   115  	if err != nil {
   116  		return nil, http.StatusRequestEntityTooLarge, err
   117  	}
   118  	return bytes, http.StatusOK, nil
   119  }