Fix: proxy routing

devopsleague · Jul 16, 2024 · 884e76a · 884e76a
1 parent 798e902
commit 884e76a
Show file tree

Hide file tree

Showing 4 changed files with 181 additions and 3 deletions.
diff --git a/internal/coordinator/api/api.go b/internal/coordinator/api/api.go
@@ -688,8 +688,11 @@ func (a *API) proxy(c echo.Context) error {
 	}
 	proxy := httputil.NewSingleHostReverseProxy(backendURL)
 	req := c.Request()
-	path := strings.Join(strings.Split(c.Path(), "/")[5:], "/")
-	req.URL.Path = fmt.Sprintf("/tasks/%s/%s/proxy/%s", task.ID, port, path)
+	path := strings.Join(strings.Split(req.URL.Path, "/")[5:], "/")
+	if path != "" && !strings.HasPrefix(path, "/") {
+		path = "/" + path
+	}
+	req.URL.Path = fmt.Sprintf("/tasks/%s/%s/proxy%s", task.ID, port, path)
 	proxy.ServeHTTP(c.Response(), req)
 	return nil
 }

diff --git a/internal/coordinator/api/api_test.go b/internal/coordinator/api/api_test.go
@@ -8,6 +8,8 @@ import (
 	"math/rand"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -962,3 +964,93 @@ func Test_proxyTaskNotRunning(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, http.StatusBadGateway, w.Code)
 }
+
+func Test_proxyTaskRoot(t *testing.T) {
+	ds := inmemory.NewInMemoryDatastore()
+
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/tasks/1234/8080/proxy", r.URL.Path)
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	svrURL, err := url.Parse(svr.URL)
+	assert.NoError(t, err)
+
+	port, err := strconv.Atoi(svrURL.Port())
+	assert.NoError(t, err)
+
+	node := &tork.Node{
+		ID:              "1234",
+		LastHeartbeatAt: time.Now().UTC(),
+		Hostname:        svrURL.Hostname(),
+		Port:            port,
+	}
+	err = ds.CreateNode(context.Background(), node)
+	assert.NoError(t, err)
+	ta := tork.Task{
+		ID:     "1234",
+		Name:   "test task",
+		State:  tork.TaskStateRunning,
+		NodeID: node.ID,
+	}
+	err = ds.CreateTask(context.Background(), &ta)
+	assert.NoError(t, err)
+	api, err := NewAPI(Config{
+		DataStore: ds,
+		Broker:    mq.NewInMemoryBroker(),
+	})
+	assert.NoError(t, err)
+	assert.NotNil(t, api)
+	req, err := http.NewRequest("GET", "/tasks/1234/proxy/8080", nil)
+	assert.NoError(t, err)
+	w := httptest.NewRecorder()
+	api.server.Handler.ServeHTTP(w, req)
+	assert.NoError(t, err)
+	assert.Equal(t, http.StatusOK, w.Code)
+}
+
+func Test_proxyTaskSomePath(t *testing.T) {
+	ds := inmemory.NewInMemoryDatastore()
+
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/tasks/1234/8080/proxy/some/path", r.URL.Path)
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	svrURL, err := url.Parse(svr.URL)
+	assert.NoError(t, err)
+
+	fmt.Println(svrURL)
+
+	port, err := strconv.Atoi(svrURL.Port())
+	assert.NoError(t, err)
+
+	node := &tork.Node{
+		ID:              "1234",
+		LastHeartbeatAt: time.Now().UTC(),
+		Hostname:        svrURL.Hostname(),
+		Port:            port,
+	}
+	err = ds.CreateNode(context.Background(), node)
+	assert.NoError(t, err)
+	ta := tork.Task{
+		ID:     "1234",
+		Name:   "test task",
+		State:  tork.TaskStateRunning,
+		NodeID: node.ID,
+	}
+	err = ds.CreateTask(context.Background(), &ta)
+	assert.NoError(t, err)
+	api, err := NewAPI(Config{
+		DataStore: ds,
+		Broker:    mq.NewInMemoryBroker(),
+	})
+	assert.NoError(t, err)
+	assert.NotNil(t, api)
+	req, err := http.NewRequest("GET", "/tasks/1234/proxy/8080/some/path", nil)
+	assert.NoError(t, err)
+	w := httptest.NewRecorder()
+	api.server.Handler.ServeHTTP(w, req)
+	assert.NoError(t, err)
+	assert.Equal(t, http.StatusOK, w.Code)
+}
diff --git a/internal/worker/api.go b/internal/worker/api.go
@@ -82,7 +82,11 @@ func (s *api) proxy(c echo.Context) error {
 	}
 	proxy := httputil.NewSingleHostReverseProxy(backendURL)
 	req := c.Request()
-	req.URL.Path = strings.Join(strings.Split(c.Path(), "/")[4:], "/")
+	path := strings.Join(strings.Split(req.URL.Path, "/")[4:], "/")
+	if path != "" && !strings.HasPrefix(path, "/") {
+		path = "/" + path
+	}
+	req.URL.Path = path
 	proxy.ServeHTTP(c.Response(), req)
 	return nil
 }

diff --git a/internal/worker/api_test.go b/internal/worker/api_test.go
@@ -4,8 +4,11 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
+	"strconv"
 	"testing"
 
+	"github.com/runabol/tork"
 	"github.com/runabol/tork/internal/syncx"
 	"github.com/runabol/tork/mq"
 	"github.com/runabol/tork/runtime/docker"
@@ -30,3 +33,79 @@ func Test_health(t *testing.T) {
 	assert.Contains(t, string(body), "\"status\":\"UP\"")
 	assert.Equal(t, http.StatusOK, w.Code)
 }
+
+func Test_proxyTaskRoot(t *testing.T) {
+	rt, err := docker.NewDockerRuntime()
+	assert.NoError(t, err)
+
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/", r.URL.Path)
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	svrURL, err := url.Parse(svr.URL)
+	assert.NoError(t, err)
+
+	port, err := strconv.Atoi(svrURL.Port())
+	assert.NoError(t, err)
+
+	tasks := &syncx.Map[string, runningTask]{}
+	tasks.Set("1234", runningTask{
+		task: &tork.Task{
+			ID: "1234",
+			Ports: []*tork.Port{{
+				Port:     "8080",
+				HostPort: port,
+			}},
+		},
+	})
+
+	api := newAPI(Config{
+		Broker:  mq.NewInMemoryBroker(),
+		Runtime: rt,
+	}, tasks)
+	assert.NotNil(t, api)
+	req, err := http.NewRequest("GET", "/tasks/1234/8080", nil)
+	assert.NoError(t, err)
+	w := httptest.NewRecorder()
+	api.server.Handler.ServeHTTP(w, req)
+	assert.Equal(t, http.StatusOK, w.Code)
+}
+
+func Test_proxyTaskSomePath(t *testing.T) {
+	rt, err := docker.NewDockerRuntime()
+	assert.NoError(t, err)
+
+	svr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/some/path", r.URL.Path)
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	svrURL, err := url.Parse(svr.URL)
+	assert.NoError(t, err)
+
+	port, err := strconv.Atoi(svrURL.Port())
+	assert.NoError(t, err)
+
+	tasks := &syncx.Map[string, runningTask]{}
+	tasks.Set("1234", runningTask{
+		task: &tork.Task{
+			ID: "1234",
+			Ports: []*tork.Port{{
+				Port:     "8080",
+				HostPort: port,
+			}},
+		},
+	})
+
+	api := newAPI(Config{
+		Broker:  mq.NewInMemoryBroker(),
+		Runtime: rt,
+	}, tasks)
+	assert.NotNil(t, api)
+	req, err := http.NewRequest("GET", "/tasks/1234/8080/some/path", nil)
+	assert.NoError(t, err)
+	w := httptest.NewRecorder()
+	api.server.Handler.ServeHTTP(w, req)
+	assert.Equal(t, http.StatusOK, w.Code)
+}