github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/engine/daemon/logger/splunk/splunkhecmock_test.go (about)

     1  package splunk // import "github.com/docker/docker/daemon/logger/splunk"
     2  
     3  import (
     4  	"compress/gzip"
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  	"io"
     9  	"net"
    10  	"net/http"
    11  	"sync"
    12  	"testing"
    13  )
    14  
    15  func (message *splunkMessage) EventAsString() (string, error) {
    16  	if val, ok := message.Event.(string); ok {
    17  		return val, nil
    18  	}
    19  	return "", fmt.Errorf("Cannot cast Event %v to string", message.Event)
    20  }
    21  
    22  func (message *splunkMessage) EventAsMap() (map[string]interface{}, error) {
    23  	if val, ok := message.Event.(map[string]interface{}); ok {
    24  		return val, nil
    25  	}
    26  	return nil, fmt.Errorf("Cannot cast Event %v to map", message.Event)
    27  }
    28  
    29  type HTTPEventCollectorMock struct {
    30  	tcpAddr     *net.TCPAddr
    31  	tcpListener *net.TCPListener
    32  
    33  	mu                  sync.Mutex
    34  	token               string
    35  	simulateServerError bool
    36  	blockingCtx         context.Context
    37  
    38  	test *testing.T
    39  
    40  	connectionVerified bool
    41  	gzipEnabled        *bool
    42  	messages           []*splunkMessage
    43  	numOfRequests      int
    44  }
    45  
    46  func NewHTTPEventCollectorMock(t *testing.T) *HTTPEventCollectorMock {
    47  	tcpAddr := &net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: 0, Zone: ""}
    48  	tcpListener, err := net.ListenTCP("tcp", tcpAddr)
    49  	if err != nil {
    50  		t.Fatal(err)
    51  	}
    52  	return &HTTPEventCollectorMock{
    53  		tcpAddr:             tcpAddr,
    54  		tcpListener:         tcpListener,
    55  		token:               "4642492F-D8BD-47F1-A005-0C08AE4657DF",
    56  		simulateServerError: false,
    57  		test:                t,
    58  		connectionVerified:  false}
    59  }
    60  
    61  func (hec *HTTPEventCollectorMock) simulateErr(b bool) {
    62  	hec.mu.Lock()
    63  	hec.simulateServerError = b
    64  	hec.mu.Unlock()
    65  }
    66  
    67  func (hec *HTTPEventCollectorMock) withBlock(ctx context.Context) {
    68  	hec.mu.Lock()
    69  	hec.blockingCtx = ctx
    70  	hec.mu.Unlock()
    71  }
    72  
    73  func (hec *HTTPEventCollectorMock) URL() string {
    74  	return "http://" + hec.tcpListener.Addr().String()
    75  }
    76  
    77  func (hec *HTTPEventCollectorMock) Serve() error {
    78  	return http.Serve(hec.tcpListener, hec)
    79  }
    80  
    81  func (hec *HTTPEventCollectorMock) Close() error {
    82  	return hec.tcpListener.Close()
    83  }
    84  
    85  func (hec *HTTPEventCollectorMock) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
    86  	var err error
    87  
    88  	hec.numOfRequests++
    89  
    90  	hec.mu.Lock()
    91  	simErr := hec.simulateServerError
    92  	ctx := hec.blockingCtx
    93  	hec.mu.Unlock()
    94  
    95  	if ctx != nil {
    96  		<-hec.blockingCtx.Done()
    97  	}
    98  
    99  	if simErr {
   100  		if request.Body != nil {
   101  			defer request.Body.Close()
   102  		}
   103  		writer.WriteHeader(http.StatusInternalServerError)
   104  		return
   105  	}
   106  
   107  	switch request.Method {
   108  	case http.MethodOptions:
   109  		// Verify that options method is getting called only once
   110  		if hec.connectionVerified {
   111  			hec.test.Errorf("Connection should not be verified more than once. Got second request with %s method.", request.Method)
   112  		}
   113  		hec.connectionVerified = true
   114  		writer.WriteHeader(http.StatusOK)
   115  	case http.MethodPost:
   116  		// Always verify that Driver is using correct path to HEC
   117  		if request.URL.String() != "/services/collector/event/1.0" {
   118  			hec.test.Errorf("Unexpected path %v", request.URL)
   119  		}
   120  		defer request.Body.Close()
   121  
   122  		if authorization, ok := request.Header["Authorization"]; !ok || authorization[0] != ("Splunk "+hec.token) {
   123  			hec.test.Error("Authorization header is invalid.")
   124  		}
   125  
   126  		gzipEnabled := false
   127  		if contentEncoding, ok := request.Header["Content-Encoding"]; ok && contentEncoding[0] == "gzip" {
   128  			gzipEnabled = true
   129  		}
   130  
   131  		if hec.gzipEnabled == nil {
   132  			hec.gzipEnabled = &gzipEnabled
   133  		} else if gzipEnabled != *hec.gzipEnabled {
   134  			// Nothing wrong with that, but we just know that Splunk Logging Driver does not do that
   135  			hec.test.Error("Driver should not change Content Encoding.")
   136  		}
   137  
   138  		var gzipReader *gzip.Reader
   139  		var reader io.Reader
   140  		if gzipEnabled {
   141  			gzipReader, err = gzip.NewReader(request.Body)
   142  			if err != nil {
   143  				hec.test.Fatal(err)
   144  			}
   145  			reader = gzipReader
   146  		} else {
   147  			reader = request.Body
   148  		}
   149  
   150  		// Read body
   151  		var body []byte
   152  		body, err = io.ReadAll(reader)
   153  		if err != nil {
   154  			hec.test.Fatal(err)
   155  		}
   156  
   157  		// Parse message
   158  		messageStart := 0
   159  		for i := 0; i < len(body); i++ {
   160  			if i == len(body)-1 || (body[i] == '}' && body[i+1] == '{') {
   161  				var message splunkMessage
   162  				err = json.Unmarshal(body[messageStart:i+1], &message)
   163  				if err != nil {
   164  					hec.test.Log(string(body[messageStart : i+1]))
   165  					hec.test.Fatal(err)
   166  				}
   167  				hec.messages = append(hec.messages, &message)
   168  				messageStart = i + 1
   169  			}
   170  		}
   171  
   172  		if gzipEnabled {
   173  			gzipReader.Close()
   174  		}
   175  
   176  		writer.WriteHeader(http.StatusOK)
   177  	default:
   178  		hec.test.Errorf("Unexpected HTTP method %s", http.MethodOptions)
   179  		writer.WriteHeader(http.StatusBadRequest)
   180  	}
   181  }