traefik/pkg/pilot/pilot_test.go

226 lines
5.2 KiB
Go
Raw Normal View History

package pilot
import (
"context"
"encoding/json"
"fmt"
"hash/fnv"
"net/http"
"net/http/httptest"
"reflect"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/traefik/traefik/v2/pkg/config/dynamic"
"github.com/traefik/traefik/v2/pkg/metrics"
"github.com/traefik/traefik/v2/pkg/safe"
)
func TestTick(t *testing.T) {
receivedConfig := make(chan bool)
mux := http.NewServeMux()
server := httptest.NewServer(mux)
t.Cleanup(server.Close)
mux.HandleFunc("/", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
err := json.NewEncoder(rw).Encode(instanceInfo{ID: "123"})
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
})
mux.HandleFunc("/command", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
receivedConfig <- true
})
pilot := New("token", metrics.RegisterPilot(), safe.NewPool(context.Background()))
pilot.client.baseInstanceInfoURL = server.URL
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
go pilot.Tick(ctx)
pilot.SetDynamicConfiguration(dynamic.Configuration{})
pilot.SetDynamicConfiguration(dynamic.Configuration{})
select {
case <-time.Tick(10 * time.Second):
t.Fatal("Timeout")
case <-receivedConfig:
return
}
}
func TestClient_SendInstanceInfo(t *testing.T) {
myToken := "myToken"
myTokenHash, err := hashToken(myToken)
require.NoError(t, err)
mux := http.NewServeMux()
server := httptest.NewServer(mux)
t.Cleanup(server.Close)
mux.HandleFunc("/", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
tk := req.Header.Get(tokenHeader)
if tk != myToken {
http.Error(rw, fmt.Sprintf("invalid token: %s", tk), http.StatusUnauthorized)
return
}
err := json.NewEncoder(rw).Encode(instanceInfo{ID: "123"})
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
})
mux.HandleFunc("/command", func(rw http.ResponseWriter, req *http.Request) {
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
tk := req.Header.Get(tokenHeader)
if tk != myToken {
http.Error(rw, fmt.Sprintf("invalid token: %s", tk), http.StatusUnauthorized)
return
}
tkh := req.Header.Get(tokenHashHeader)
if tkh != myTokenHash {
http.Error(rw, fmt.Sprintf("invalid token hash: %s", tkh), http.StatusBadRequest)
return
}
defer func() { _ = req.Body.Close() }()
info := &instanceInfo{}
err := json.NewDecoder(req.Body).Decode(info)
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
if info.ID != "123" {
http.Error(rw, fmt.Sprintf("invalid ID: %s", info.ID), http.StatusBadRequest)
}
})
client := client{
baseInstanceInfoURL: server.URL,
httpClient: http.DefaultClient,
token: myToken,
tokenHash: myTokenHash,
}
err = client.SendInstanceInfo(context.Background(), []metrics.PilotMetric{})
require.NoError(t, err)
}
func TestClient_SendAnonDynConf(t *testing.T) {
myToken := "myToken"
myTokenHash, err := hashToken(myToken)
require.NoError(t, err)
var count int
mux := http.NewServeMux()
mux.HandleFunc("/collect", func(rw http.ResponseWriter, req *http.Request) {
count++
if count == 1 {
http.Error(rw, "OOPS", http.StatusBadRequest)
return
}
if req.Method != http.MethodPost {
http.Error(rw, "invalid method", http.StatusMethodNotAllowed)
return
}
tkh := req.Header.Get(tokenHashHeader)
if tkh != myTokenHash {
http.Error(rw, fmt.Sprintf("invalid token hash: %s", tkh), http.StatusBadRequest)
return
}
defer func() { _ = req.Body.Close() }()
config := &dynamic.Configuration{}
err := json.NewDecoder(req.Body).Decode(config)
if err != nil {
http.Error(rw, err.Error(), http.StatusInternalServerError)
return
}
router, exists := config.HTTP.Routers["foo"]
if !exists {
http.Error(rw, "router configuration is missing", http.StatusBadRequest)
return
}
if !reflect.DeepEqual(router, &dynamic.Router{Service: "foo", Rule: "xxxx"}) {
http.Error(rw, fmt.Sprintf("configuration is not anonymized: %+v", router), http.StatusBadRequest)
return
}
})
server := httptest.NewServer(mux)
t.Cleanup(server.Close)
client := client{
baseGatewayURL: server.URL,
httpClient: http.DefaultClient,
token: myToken,
tokenHash: myTokenHash,
}
config := dynamic.Configuration{
HTTP: &dynamic.HTTPConfiguration{
Routers: map[string]*dynamic.Router{
"foo": {
Service: "foo",
Rule: "foo.com",
},
},
},
}
err = client.SendAnonDynConf(context.Background(), config)
require.NoError(t, err)
assert.Equal(t, 2, count)
}
func hashToken(token string) (string, error) {
tokenHash := fnv.New64a()
_, err := tokenHash.Write([]byte(token))
if err != nil {
return "", err
}
return fmt.Sprintf("%x", tokenHash.Sum64()), nil
}