Increase jwt-provider test coverage

services/jwt-provider/handler_create_jwt.go

@@ -1,10 +1,15 @@
 package jwt_provider
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"net"
 	"net/http"
+	"strings"
+
+	"github.com/docker/docker/api/types"
+	"github.com/forta-network/forta-node/config"
 )
 
 type CreateJWTMessage struct {
@@ -14,11 +19,62 @@ type CreateJWTResponse struct {
 	Token string `json:"token"`
 }
 
+const envPrefix = config.EnvFortaBotID + "="
+
 const (
 	errBadCreateMessage  = "bad create jwt message body"
 	errFailedToCreateJWT = "can't find bot id from request source"
 )
 
+// agentIDReverseLookup reverse lookup from ip to agent id.
+func (j *JWTProvider) agentIDReverseLookup(ctx context.Context, ipAddr string) (string, error) {
+	container, err := j.findContainerByIP(ctx, ipAddr)
+	if err != nil {
+		return "", err
+	}
+
+	botID, err := j.extractBotIDFromContainer(ctx, container)
+	if err != nil {
+		return "", err
+	}
+
+	return botID, nil
+}
+
+func (j *JWTProvider) extractBotIDFromContainer(ctx context.Context, container types.Container) (string, error) {
+	// container struct doesn't have the "env" information, inspection required.
+	c, err := j.dockerClient.InspectContainer(ctx, container.ID)
+	if err != nil {
+		return "", err
+	}
+
+	// find the env variable with bot id
+	for _, s := range c.Config.Env {
+		if env := strings.SplitAfter(s, envPrefix); len(env) == 2 {
+			return env[1], nil
+		}
+	}
+
+	return "", fmt.Errorf("can't extract bot id from container")
+}
+
+func (j *JWTProvider) findContainerByIP(ctx context.Context, ipAddr string) (types.Container, error) {
+	containers, err := j.dockerClient.GetContainers(ctx)
+	if err != nil {
+		return types.Container{}, err
+	}
+
+	// find the container that has the same ip
+	for _, container := range containers {
+		for _, network := range container.NetworkSettings.Networks {
+			if network.IPAddress == ipAddr {
+				return container, nil
+			}
+		}
+	}
+	return types.Container{}, fmt.Errorf("can't find container %s", ipAddr)
+}
+
 // createJWTHandler returns a scanner jwt token with claims [hash] = hash(uri,payload) and [bot] = "bot id"
 func (j *JWTProvider) createJWTHandler(w http.ResponseWriter, req *http.Request) {
 	var msg CreateJWTMessage
@@ -31,29 +87,34 @@ func (j *JWTProvider) createJWTHandler(w http.ResponseWriter, req *http.Request)
 		}
 	}
 
-	ipAddr, _, err := net.SplitHostPort(req.RemoteAddr)
+	jwt, err := j.doCreateJWT(req.Context(), req.RemoteAddr, msg.Claims)
 	if err != nil {
 		w.WriteHeader(http.StatusBadRequest)
-		_, _ = fmt.Fprintf(w, "can't extract ip from request %s", req.RemoteAddr)
+		_, _ = fmt.Fprintf(w, "can not create jwt: %v", err)
 		return
 	}
 
-	agentID, err := j.agentIDReverseLookup(req.Context(), ipAddr)
+	resp, err := json.Marshal(CreateJWTResponse{Token: jwt})
+
+	w.WriteHeader(http.StatusOK)
+	_, _ = fmt.Fprintf(w, "%s", resp)
+}
+
+func (j *JWTProvider) doCreateJWT(ctx context.Context, remoteAddr string, claims map[string]interface{}) (string, error) {
+	ipAddr, _, err := net.SplitHostPort(remoteAddr)
 	if err != nil {
-		w.WriteHeader(http.StatusBadRequest)
-		_, _ = fmt.Fprintf(w, "can't find bot id from request source %s, err: %v", ipAddr, err)
-		return
+		return "", fmt.Errorf("can't extract ip from request %s", remoteAddr)
 	}
 
-	jwt, err := CreateBotJWT(j.cfg.Key, agentID, msg.Claims)
+	agentID, err := j.agentIDReverseLookup(ctx, ipAddr)
 	if err != nil {
-		w.WriteHeader(http.StatusBadRequest)
-		_, _ = fmt.Fprint(w, errFailedToCreateJWT)
-		return
+		return "", fmt.Errorf("can't find bot id from request source %s, err: %v", ipAddr, err)
 	}
 
-	resp, err := json.Marshal(CreateJWTResponse{Token: jwt})
+	jwt, err := CreateBotJWT(j.cfg.Key, agentID, claims)
+	if err != nil {
+		return "", fmt.Errorf("%s: %v", errFailedToCreateJWT, err)
+	}
 
-	w.WriteHeader(http.StatusOK)
-	_, _ = fmt.Fprintf(w, "%s", resp)
-}
+	return jwt, nil
+}

services/jwt-provider/handler_create_jwt_test.go

@@ -0,0 +1,216 @@
+package jwt_provider
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/network"
+	"github.com/ethereum/go-ethereum/accounts/keystore"
+	"github.com/forta-network/forta-core-go/security"
+	mock_clients "github.com/forta-network/forta-node/clients/mocks"
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+)
+
+func mockJWTProvider(t *testing.T) (*JWTProvider, *mock_clients.MockDockerClient) {
+	ctrl := gomock.NewController(t)
+
+	mockDockerClient := mock_clients.NewMockDockerClient(ctrl)
+
+	dir := t.TempDir()
+	ks := keystore.NewKeyStore(dir, keystore.StandardScryptN, keystore.StandardScryptP)
+
+	_, err := ks.NewAccount("Forta123")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	mockKey, err := security.LoadKeyWithPassphrase(dir, "Forta123")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return &JWTProvider{
+		cfg:          &JWTProviderConfig{Key: mockKey},
+		dockerClient: mockDockerClient,
+	}, mockDockerClient
+
+}
+
+func TestJWTProvider_jwtHandler(t *testing.T) {
+	j, mockDockerClient := mockJWTProvider(t)
+
+	// test context
+	ctx := context.Background()
+
+	//
+	// Test Case 1: Bot can retrieve JWT
+	//
+	mockBotIP := "172.0.1.2"
+	mockBotRemoteAddr := mockBotIP + ":8423"
+	mockBotID := "1"
+	mockBotContainerInfo := []types.Container{
+		{
+			NetworkSettings: &types.SummaryNetworkSettings{
+				Networks: map[string]*network.EndpointSettings{
+					"local": {IPAddress: mockBotIP},
+				},
+			},
+		},
+	}
+	mockBotContainerJSON := &types.ContainerJSON{
+		Config: &container.Config{
+			Env: []string{
+				envPrefix + mockBotID,
+			},
+		},
+	}
+
+	mockDockerClient.EXPECT().GetContainers(gomock.Any()).Return(mockBotContainerInfo, nil)
+	mockDockerClient.EXPECT().InspectContainer(gomock.Any(), gomock.Any()).Return(mockBotContainerJSON, nil)
+
+	_, err := j.doCreateJWT(ctx, mockBotRemoteAddr, nil)
+	assert.NoError(t, err)
+
+	//
+	// Test Case 2: Unknown sources can't get JWT
+	//
+	mockBotIP = "172.0.1.2"
+	mockBotRemoteAddr = mockBotIP + ":8423"
+	mockBotID = "1"
+	mockBotContainerInfo = []types.Container{
+		{
+			NetworkSettings: &types.SummaryNetworkSettings{
+				Networks: map[string]*network.EndpointSettings{},
+			},
+		},
+	}
+	mockBotContainerJSON = &types.ContainerJSON{
+		Config: &container.Config{
+			Env: []string{
+				envPrefix + mockBotID,
+			},
+		},
+	}
+
+	mockDockerClient.EXPECT().GetContainers(gomock.Any()).Return(mockBotContainerInfo, nil)
+
+	_, err = j.doCreateJWT(ctx, mockBotRemoteAddr, nil)
+	assert.Error(t, err)
+
+	//
+	// Test Case 3: Source with bad remote address
+	//
+	mockBotIP = "www.X.Y.Z"
+	mockBotRemoteAddr = mockBotIP
+	mockBotID = "1"
+	mockBotContainerInfo = []types.Container{
+		{
+			NetworkSettings: &types.SummaryNetworkSettings{
+				Networks: map[string]*network.EndpointSettings{},
+			},
+		},
+	}
+	mockBotContainerJSON = &types.ContainerJSON{
+		Config: &container.Config{
+			Env: []string{
+				envPrefix + mockBotID,
+			},
+		},
+	}
+
+	_, err = j.doCreateJWT(ctx, mockBotRemoteAddr, nil)
+	assert.Error(t, err)
+
+	//
+	// Test Case 4: Source with bad remote address
+	//
+	mockBotIP = "172.0.1.2"
+	mockBotRemoteAddr = mockBotIP + ":8423"
+	mockBotID = "1"
+	mockBotContainerInfo = []types.Container{
+		{
+			NetworkSettings: &types.SummaryNetworkSettings{
+				Networks: map[string]*network.EndpointSettings{
+					"local": {IPAddress: mockBotIP},
+				},
+			},
+		},
+	}
+	mockBotContainerJSON = &types.ContainerJSON{
+		Config: &container.Config{
+			Env: []string{
+				envPrefix + mockBotID,
+			},
+		},
+	}
+
+	j2, mockDockerClient2 := mockJWTProvider(t)
+	j2.cfg.Key = nil
+	mockDockerClient2.EXPECT().GetContainers(gomock.Any()).Return(mockBotContainerInfo, nil)
+	mockDockerClient2.EXPECT().InspectContainer(gomock.Any(), gomock.Any()).Return(mockBotContainerJSON, nil)
+
+	_, err = j2.doCreateJWT(ctx, mockBotRemoteAddr, nil)
+	assert.Error(t, err)
+
+}
+
+func TestJWTProvider_createJWTHandler(t *testing.T) {
+	j, mockDockerClient := mockJWTProvider(t)
+
+	body := CreateJWTMessage{map[string]interface{}{"test-claim": "success"}}
+	b := bytes.NewBuffer([]byte{})
+	_ = json.NewEncoder(b).Encode(body)
+	req := httptest.NewRequest(http.MethodPost, "/create", b)
+	w := httptest.NewRecorder()
+
+	mockBotID := "1"
+	mockBotIP, _, _ := net.SplitHostPort(req.RemoteAddr)
+	mockBotContainerInfo := []types.Container{
+		{
+			NetworkSettings: &types.SummaryNetworkSettings{
+				Networks: map[string]*network.EndpointSettings{
+					"local": {
+						IPAddress: mockBotIP,
+					},
+				},
+			},
+		},
+	}
+	mockBotContainerJSON := &types.ContainerJSON{
+		Config: &container.Config{
+			Env: []string{
+				envPrefix + mockBotID,
+			},
+		},
+	}
+
+	// Test case 1: can retrieve token
+	mockDockerClient.EXPECT().GetContainers(gomock.Any()).Return(mockBotContainerInfo, nil)
+	mockDockerClient.EXPECT().InspectContainer(gomock.Any(), gomock.Any()).Return(mockBotContainerJSON, nil)
+	j.createJWTHandler(w, req)
+
+	// Test case 2: bad request body
+	b2 := bytes.NewBuffer([]byte("xxxxxx"))
+	req2 := httptest.NewRequest(http.MethodPost, "/create", b2)
+	j.createJWTHandler(w, req2)
+
+	// Test case 3: can not retrieve token, bad source
+	body3 := CreateJWTMessage{map[string]interface{}{"test-claim": "success"}}
+	b3 := bytes.NewBuffer([]byte{})
+	_ = json.NewEncoder(b3).Encode(body3)
+	req3 := httptest.NewRequest(http.MethodPost, "/create", b3)
+	req3.RemoteAddr = "1.1.1.1:4444"
+	w3 := httptest.NewRecorder()
+
+	mockDockerClient.EXPECT().GetContainers(gomock.Any()).Return(mockBotContainerInfo, nil)
+
+	j.createJWTHandler(w3, req3)
+}