package store import ( "bytes" "errors" "io" "mime" "net/http" "net/url" "strings" "sync" "time" "github.com/google/uuid" ) const ( defaultQueueSize = 100 maxQueueSize = 100 ) var ErrRecordNotFound = errors.New("record not found") type TunnelStore interface { Get(before time.Time, count int) (results []*TunnelRecord, more bool) RecordTCP() RecordRequest(req *http.Request) RecordResponse(resp *http.Response) error } type TunnelRecord struct { ID uuid.UUID Time time.Time URL *url.URL Method string Status int RequestHeaders http.Header RequestBodyType string RequestBody []byte ResponseHeaders http.Header ResponseBodyType string ResponseBody []byte } func NewTunnelStore(queueSize int) TunnelStore { if queueSize <= 0 { queueSize = defaultQueueSize } else if queueSize > maxQueueSize { queueSize = maxQueueSize } return &tunnelStoreImpl{queueSize: queueSize} } type tunnelStoreImpl struct { orderedRecords []*TunnelRecord queueSize int mu sync.Mutex } func (s *tunnelStoreImpl) Get(before time.Time, count int) ([]*TunnelRecord, bool) { // Find First start := -1 for i, r := range s.orderedRecords { if r.Time.Before(before) { start = i break } } // Not Found if start == -1 { return nil, false } // Subslice Records end := min(start+count, len(s.orderedRecords)) results := s.orderedRecords[start:end] more := end < len(s.orderedRecords) return results, more } func (s *tunnelStoreImpl) RecordRequest(req *http.Request) { s.mu.Lock() defer s.mu.Unlock() url := *req.URL rec := &TunnelRecord{ ID: uuid.New(), Time: time.Now(), URL: &url, Method: req.Method, RequestHeaders: req.Header, RequestBodyType: req.Header.Get("Content-Type"), } if bodyData, err := getRequestBody(req); err == nil { rec.RequestBody = bodyData } // Add Record & Truncate s.orderedRecords = append(s.orderedRecords, rec) if len(s.orderedRecords) > s.queueSize { s.orderedRecords = s.orderedRecords[len(s.orderedRecords)-s.queueSize:] } *req = *req.WithContext(withRecord(req.Context(), rec)) } func (s *tunnelStoreImpl) RecordResponse(resp *http.Response) error { rec, found := getRecord(resp.Request.Context()) if !found { return ErrRecordNotFound } rec.ResponseHeaders = resp.Header rec.ResponseBodyType = resp.Header.Get("Content-Type") if bodyData, err := getResponseBody(resp); err == nil { rec.ResponseBody = bodyData } return nil } func (s *tunnelStoreImpl) RecordTCP() { s.mu.Lock() defer s.mu.Unlock() // TODO } func getRequestBody(req *http.Request) ([]byte, error) { if req.ContentLength == 0 || req.Body == nil || req.Body == http.NoBody { return nil, nil } if !isTextContentType(req.Header.Get("Content-Type")) { return nil, nil } // Read Body bodyBytes, err := io.ReadAll(req.Body) if err != nil { return nil, err } // Restore Body req.Body = io.NopCloser(bytes.NewReader(bodyBytes)) return bodyBytes, nil } func getResponseBody(resp *http.Response) ([]byte, error) { if resp.ContentLength == 0 || resp.Body == nil || resp.Body == http.NoBody { return nil, nil } if !isTextContentType(resp.Header.Get("Content-Type")) { return nil, nil } // Read Body bodyBytes, err := io.ReadAll(resp.Body) if err != nil { return nil, err } // Restore Body resp.Body = io.NopCloser(bytes.NewReader(bodyBytes)) return bodyBytes, nil } func isTextContentType(contentType string) bool { mediaType, _, err := mime.ParseMediaType(contentType) if err != nil { return false } if strings.HasPrefix(mediaType, "text/") { return true } switch mediaType { case "application/json": return true case "application/xml": return true case "application/x-www-form-urlencoded": return true default: return false } }