diff --git a/cmd/rdpgw/web/web_test.go b/cmd/rdpgw/web/web_test.go index 166d3cb0d4f4799d215e90a5dd59a48aa528ce07..6ceccdf7702c97519627f2c0cf5f09b35720ab12 100644 --- a/cmd/rdpgw/web/web_test.go +++ b/cmd/rdpgw/web/web_test.go @@ -3,10 +3,18 @@ package web import ( "context" "github.com/bolkedebruin/rdpgw/cmd/rdpgw/security" + "net/http" + "net/http/httptest" "net/url" + "strings" "testing" ) +const ( + testuser = "test_user" + gateway = "https://my.gateway.com:993" +) + var ( hosts = []string{"10.0.0.1:3389", "10.1.1.1:3000", "32.32.11.1", "remote.host.com"} key = []byte("thisisasessionkeyreplacethisjetzt") @@ -97,3 +105,79 @@ func TestGetHost(t *testing.T) { t.Fatalf("%s does not equal %s", host, hosts[0]) } } + +func TestHandler_HandleDownload(t *testing.T) { + req, err := http.NewRequest("GET", "/connect", nil) + if err != nil { + t.Fatal(err) + } + + rr := httptest.NewRecorder() + ctx := req.Context() + ctx = context.WithValue(ctx, "preferred_username", testuser) + req = req.WithContext(ctx) + + u, _ := url.Parse(gateway) + c := Config{ + HostSelection: "roundrobin", + Hosts: hosts, + PAATokenGenerator: paaTokenMock, + GatewayAddress: u, + RdpOpts: RdpOpts{SplitUserDomain: true}, + } + h := c.NewHandler() + + hh := http.HandlerFunc(h.HandleDownload) + hh.ServeHTTP(rr, req) + + if status := rr.Code; status != http.StatusOK { + t.Errorf("handler returned wrong status code: got %v want %v", + status, http.StatusOK) + } + + if ctype := rr.Header().Get("Content-Type"); ctype != "application/x-rdp" { + t.Errorf("content type header does not match: got %v want %v", + ctype, "application/json") + } + + if cdisp := rr.Header().Get("Content-Disposition"); cdisp == "" { + t.Errorf("content disposition is nil") + } + + data := rdpToMap(strings.Split(rr.Body.String(), crlf)) + if data["username"] != testuser { + t.Errorf("username key in rdp does not match: got %v want %v", data["username"], testuser) + } + + if data["gatewayhostname"] != u.Host { + t.Errorf("gatewayhostname key in rdp does not match: got %v want %v", data["gatewayhostname"], u.Host) + } + + if token, _ := paaTokenMock(ctx, testuser, data["full address"]); token != data["gatewayaccesstoken"] { + t.Errorf("gatewayaccesstoken key in rdp does not match username_full address: got %v want %v", + data["gatewayaccesstoken"], token) + } + + if !contains(data["full address"], hosts) { + t.Errorf("full address key in rdp is not in allowed hosts list: go %v want in %v", + data["full address"], hosts) + } + +} + +func paaTokenMock(ctx context.Context, username string, host string) (string, error) { + return username + "_" + host, nil +} + +func rdpToMap(rdp []string) map[string]string { + ret := make(map[string]string) + + for s := range rdp { + d := strings.SplitN(rdp[s], ":", 3) + if len(d) >= 2 { + ret[d[0]] = d[2] + } + } + + return ret +}