[Unit Tests] - PubKeyContext

[tomsmith8]

Unit Test Coverage for "PubKeyContext"

---

Stakwork Run

---

Unit Test Code

---

package auth

import (
  "context"
  "errors"
  "net/http"
  "net/http/httptest"
  "testing"
  "time"

  "github.com/dgrijalva/jwt-go"
  "github.com/stretchr/testify/assert"
  "github.com/stretchr/testify/mock"
)

// Mocking DecodeJwt function
type MockDecoder struct {
  mock.Mock
}

func (m *MockDecoder) DecodeJwt(token string) (jwt.MapClaims, error) {
  args := m.Called(token)
  return args.Get(0).(jwt.MapClaims), args.Error(1)
}

// Mocking VerifyTribeUUID function
type MockVerifier struct {
  mock.Mock
}

func (m *MockVerifier) VerifyTribeUUID(uuid string, checkTimestamp bool) (string, error) {
  args := m.Called(uuid, checkTimestamp)
  return args.String(0), args.Error(1)
}

func TestPubKeyContext(t *testing.T) {
  mockDecoder := new(MockDecoder)
  mockVerifier := new(MockVerifier)

  // Define a next handler to capture if it was called
  nextHandlerCalled := false
  nextHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  	nextHandlerCalled = true
  })

  // Test cases
  tests := []struct {
  	name           string
  	token          string
  	header         bool
  	mockDecode     func()
  	mockVerify     func()
  	expectedStatus int
  	expectNextCall bool
  }{
  	{
  		name:   "Valid JWT Token in Query Parameter",
  		token:  "valid.jwt.token",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "valid.jwt.token").Return(jwt.MapClaims{"pubkey": "valid_pubkey"}, nil)
  		},
  		expectedStatus: http.StatusOK,
  		expectNextCall: true,
  	},
  	{
  		name:   "Valid JWT Token in Header",
  		token:  "valid.jwt.token",
  		header: true,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "valid.jwt.token").Return(jwt.MapClaims{"pubkey": "valid_pubkey"}, nil)
  		},
  		expectedStatus: http.StatusOK,
  		expectNextCall: true,
  	},
  	{
  		name:   "Valid Tribe UUID Token",
  		token:  "valid_tribe_uuid",
  		header: false,
  		mockVerify: func() {
  			mockVerifier.On("VerifyTribeUUID", "valid_tribe_uuid", true).Return("valid_pubkey", nil)
  		},
  		expectedStatus: http.StatusOK,
  		expectNextCall: true,
  	},
  	{
  		name:           "Empty Token in Query and Header",
  		token:          "",
  		header:         false,
  		expectedStatus: http.StatusUnauthorized,
  		expectNextCall: false,
  	},
  	{
  		name:   "Expired JWT Token",
  		token:  "expired.jwt.token",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "expired.jwt.token").Return(jwt.MapClaims{"exp": time.Now().Add(-time.Hour).Unix()}, nil)
  		},
  		expectedStatus: http.StatusUnauthorized,
  		expectNextCall: false,
  	},
  	{
  		name:   "JWT Token with No pubkey Claim",
  		token:  "no_pubkey.jwt.token",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "no_pubkey.jwt.token").Return(jwt.MapClaims{}, nil)
  		},
  		expectedStatus: http.StatusUnauthorized,
  		expectNextCall: false,
  	},
  	{
  		name:   "JWT Token with Additional Unused Claims",
  		token:  "unused_claims.jwt.token",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "unused_claims.jwt.token").Return(jwt.MapClaims{"pubkey": "valid_pubkey", "extra": "value"}, nil)
  		},
  		expectedStatus: http.StatusOK,
  		expectNextCall: true,
  	},
  	{
  		name:   "Invalid JWT Token Format",
  		token:  "invalid.jwt.token",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "invalid.jwt.token").Return(nil, errors.New("invalid token"))
  		},
  		expectedStatus: http.StatusUnauthorized,
  		expectNextCall: false,
  	},
  	{
  		name:   "Invalid Tribe UUID Token",
  		token:  "invalid_tribe_uuid",
  		header: false,
  		mockVerify: func() {
  			mockVerifier.On("VerifyTribeUUID", "invalid_tribe_uuid", true).Return("", errors.New("invalid token"))
  		},
  		expectedStatus: http.StatusUnauthorized,
  		expectNextCall: false,
  	},
  	{
  		name:   "Error in JWT Decoding",
  		token:  "error_decoding.jwt.token",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "error_decoding.jwt.token").Return(nil, errors.New("decoding error"))
  		},
  		expectedStatus: http.StatusUnauthorized,
  		expectNextCall: false,
  	},
  	{
  		name:   "Error in Tribe UUID Verification",
  		token:  "error_verification_tribe_uuid",
  		header: false,
  		mockVerify: func() {
  			mockVerifier.On("VerifyTribeUUID", "error_verification_tribe_uuid", true).Return("", errors.New("verification error"))
  		},
  		expectedStatus: http.StatusUnauthorized,
  		expectNextCall: false,
  	},
  	{
  		name:   "Token with Special Characters",
  		token:  "special!@#$%^&*()_+",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "special!@#$%^&*()_+").Return(jwt.MapClaims{"pubkey": "valid_pubkey"}, nil)
  		},
  		expectedStatus: http.StatusOK,
  		expectNextCall: true,
  	},
  	{
  		name:   "Token with Leading Dot",
  		token:  ".leading.dot.token",
  		header: false,
  		mockVerify: func() {
  			mockVerifier.On("VerifyTribeUUID", ".leading.dot.token", true).Return("valid_pubkey", nil)
  		},
  		expectedStatus: http.StatusOK,
  		expectNextCall: true,
  	},
  	{
  		name:   "Token with Different Encoding",
  		token:  "base64url_encoded_token",
  		header: false,
  		mockDecode: func() {
  			mockDecoder.On("DecodeJwt", "base64url_encoded_token").Return(jwt.MapClaims{"pubkey": "valid_pubkey"}, nil)
  		},
  		expectedStatus: http.StatusOK,
  		expectNextCall: true,
  	},
  }

  for _, tt := range tests {
  	t.Run(tt.name, func(t *testing.T) {
  		// Reset the next handler call flag
  		nextHandlerCalled = false

  		// Setup mocks
  		if tt.mockDecode != nil {
  			tt.mockDecode()
  		}
  		if tt.mockVerify != nil {
  			tt.mockVerify()
  		}

  		// Create a request
  		req := httptest.NewRequest("GET", "http://example.com", nil)
  		if tt.header {
  			req.Header.Set("x-jwt", tt.token)
  		} else {
  			q := req.URL.Query()
  			q.Add("token", tt.token)
  			req.URL.RawQuery = q.Encode()
  		}

  		// Create a response recorder
  		rr := httptest.NewRecorder()

  		// Create the handler
  		handler := PubKeyContext(nextHandler)

  		// Serve the request
  		handler.ServeHTTP(rr, req)

  		// Check the status code
  		assert.Equal(t, tt.expectedStatus, rr.Code)

  		// Check if the next handler was called
  		assert.Equal(t, tt.expectNextCall, nextHandlerCalled)

  		// Assert expectations
  		mockDecoder.AssertExpectations(t)
  		mockVerifier.AssertExpectations(t)
  	})
  }
}