github.com/LazyboyChen7/engine@v17.12.1-ce-rc2+incompatible/daemon/logger/splunk/splunkhecmock_test.go (about)

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