tls_test.go

/*
Copyright 2018-2022 Mailgun Technologies Inc

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package gubernator_test

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"strings"
	"testing"

	gubernator "github.com/mailgun/gubernator/v2"
	"github.com/mailgun/holster/v4/clock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/http2"
)

func spawnDaemon(t *testing.T, conf gubernator.DaemonConfig) *gubernator.Daemon {
	t.Helper()

	ctx, cancel := context.WithTimeout(context.Background(), clock.Second*10)
	d, err := gubernator.SpawnDaemon(ctx, conf)
	cancel()
	require.NoError(t, err)
	d.SetPeers([]gubernator.PeerInfo{{GRPCAddress: conf.GRPCListenAddress, IsOwner: true}})
	return d
}

func makeRequest(t *testing.T, conf gubernator.DaemonConfig) error {
	t.Helper()

	client, err := gubernator.DialV1Server(conf.GRPCListenAddress, conf.TLS.ClientTLS)
	require.NoError(t, err)

	resp, err := client.GetRateLimits(context.Background(), &gubernator.GetRateLimitsReq{
		Requests: []*gubernator.RateLimitReq{
			{
				Name:      "test_tls",
				UniqueKey: "account:995",
				Algorithm: gubernator.Algorithm_TOKEN_BUCKET,
				Duration:  gubernator.Second * 30,
				Limit:     100,
				Hits:      1,
			},
		},
	})

	if err != nil {
		return err
	}
	rl := resp.Responses[0]
	assert.Equal(t, "", rl.Error)
	return nil
}

func TestSetupTLS(t *testing.T) {
	tests := []struct {
		tls  *gubernator.TLSConfig
		name string
	}{
		{
			name: "user provided certificates",
			tls: &gubernator.TLSConfig{
				CaFile:   "contrib/certs/ca.cert",
				CertFile: "contrib/certs/gubernator.pem",
				KeyFile:  "contrib/certs/gubernator.key",
			},
		},
		{
			name: "user provided certificate without IP SANs",
			tls: &gubernator.TLSConfig{
				CaFile:               "contrib/certs/ca.cert",
				CertFile:             "contrib/certs/gubernator_no_ip_san.pem",
				KeyFile:              "contrib/certs/gubernator_no_ip_san.key",
				ClientAuthServerName: "gubernator",
			},
		},
		{
			name: "auto tls",
			tls: &gubernator.TLSConfig{
				AutoTLS: true,
			},
		},
		{
			name: "generate server certs with user provided ca",
			tls: &gubernator.TLSConfig{
				CaFile:    "contrib/certs/ca.cert",
				CaKeyFile: "contrib/certs/ca.key",
				AutoTLS:   true,
			},
		},
		{
			name: "client auth enabled",
			tls: &gubernator.TLSConfig{
				CaFile:     "contrib/certs/ca.cert",
				CaKeyFile:  "contrib/certs/ca.key",
				AutoTLS:    true,
				ClientAuth: tls.RequireAndVerifyClientCert,
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			conf := gubernator.DaemonConfig{
				GRPCListenAddress: "127.0.0.1:9695",
				HTTPListenAddress: "127.0.0.1:9685",
				TLS:               tt.tls,
			}

			d := spawnDaemon(t, conf)

			client, err := gubernator.DialV1Server(conf.GRPCListenAddress, tt.tls.ClientTLS)
			require.NoError(t, err)

			resp, err := client.GetRateLimits(context.Background(), &gubernator.GetRateLimitsReq{
				Requests: []*gubernator.RateLimitReq{
					{
						Name:      "test_tls",
						UniqueKey: "account:995",
						Algorithm: gubernator.Algorithm_TOKEN_BUCKET,
						Duration:  gubernator.Second * 30,
						Limit:     100,
						Hits:      1,
					},
				},
			})
			require.NoError(t, err)

			rl := resp.Responses[0]
			assert.Equal(t, "", rl.Error)
			assert.Equal(t, gubernator.Status_UNDER_LIMIT, rl.Status)
			assert.Equal(t, int64(99), rl.Remaining)
			d.Close()
		})
	}
}

func TestSetupTLSSkipVerify(t *testing.T) {
	conf := gubernator.DaemonConfig{
		GRPCListenAddress: "127.0.0.1:9695",
		HTTPListenAddress: "127.0.0.1:9685",
		TLS: &gubernator.TLSConfig{
			CaFile:   "contrib/certs/ca.cert",
			CertFile: "contrib/certs/gubernator.pem",
			KeyFile:  "contrib/certs/gubernator.key",
		},
	}

	d := spawnDaemon(t, conf)
	defer d.Close()

	tls := &gubernator.TLSConfig{
		AutoTLS:            true,
		InsecureSkipVerify: true,
	}

	err := gubernator.SetupTLS(tls)
	require.NoError(t, err)
	conf.TLS = tls

	err = makeRequest(t, conf)
	require.NoError(t, err)
}

func TestSetupTLSClientAuth(t *testing.T) {
	serverTLS := gubernator.TLSConfig{
		CaFile:           "contrib/certs/ca.cert",
		CertFile:         "contrib/certs/gubernator.pem",
		KeyFile:          "contrib/certs/gubernator.key",
		ClientAuth:       tls.RequireAndVerifyClientCert,
		ClientAuthCaFile: "contrib/certs/client-auth-ca.pem",
	}

	conf := gubernator.DaemonConfig{
		GRPCListenAddress: "127.0.0.1:9695",
		HTTPListenAddress: "127.0.0.1:9685",
		TLS:               &serverTLS,
	}

	d := spawnDaemon(t, conf)
	defer d.Close()

	// Given generated client certs
	tls := &gubernator.TLSConfig{
		AutoTLS:            true,
		InsecureSkipVerify: true,
	}

	err := gubernator.SetupTLS(tls)
	require.NoError(t, err)
	conf.TLS = tls

	// Should not be allowed without a cert signed by the client CA
	err = makeRequest(t, conf)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "code = Unavailable desc")

	// Given the client auth certs
	tls = &gubernator.TLSConfig{
		CertFile:           "contrib/certs/client-auth.pem",
		KeyFile:            "contrib/certs/client-auth.key",
		InsecureSkipVerify: true,
	}

	err = gubernator.SetupTLS(tls)
	require.NoError(t, err)
	conf.TLS = tls

	// Should be allowed to connect and make requests
	err = makeRequest(t, conf)
	require.NoError(t, err)
}

func TestTLSClusterWithClientAuthentication(t *testing.T) {
	serverTLS := gubernator.TLSConfig{
		CaFile:     "contrib/certs/ca.cert",
		CertFile:   "contrib/certs/gubernator.pem",
		KeyFile:    "contrib/certs/gubernator.key",
		ClientAuth: tls.RequireAndVerifyClientCert,
	}

	d1 := spawnDaemon(t, gubernator.DaemonConfig{
		GRPCListenAddress: "127.0.0.1:9695",
		HTTPListenAddress: "127.0.0.1:9685",
		TLS:               &serverTLS,
	})
	defer d1.Close()

	d2 := spawnDaemon(t, gubernator.DaemonConfig{
		GRPCListenAddress: "127.0.0.1:9696",
		HTTPListenAddress: "127.0.0.1:9686",
		TLS:               &serverTLS,
	})
	defer d2.Close()

	peers := []gubernator.PeerInfo{
		{
			GRPCAddress: d1.GRPCListeners[0].Addr().String(),
			HTTPAddress: d1.HTTPListener.Addr().String(),
		},
		{
			GRPCAddress: d2.GRPCListeners[0].Addr().String(),
			HTTPAddress: d2.HTTPListener.Addr().String(),
		},
	}
	d1.SetPeers(peers)
	d2.SetPeers(peers)

	// Should result in a remote call to d2
	err := makeRequest(t, d1.Config())
	require.NoError(t, err)

	config := d2.Config()
	client := &http.Client{
		Transport: &http2.Transport{
			TLSClientConfig: config.ClientTLS(),
		},
	}

	resp, err := client.Get(fmt.Sprintf("https://%s/metrics", config.HTTPListenAddress))
	require.NoError(t, err)
	defer resp.Body.Close()

	b, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	// Should have called GetPeerRateLimits on d2
	assert.Contains(t, string(b), `{method="/pb.gubernator.PeersV1/GetPeerRateLimits"} 1`)
}

func TestHTTPSClientAuth(t *testing.T) {
	conf := gubernator.DaemonConfig{
		GRPCListenAddress:       "127.0.0.1:9695",
		HTTPListenAddress:       "127.0.0.1:9685",
		HTTPStatusListenAddress: "127.0.0.1:9686",
		TLS: &gubernator.TLSConfig{
			CaFile:     "contrib/certs/ca.cert",
			CertFile:   "contrib/certs/gubernator.pem",
			KeyFile:    "contrib/certs/gubernator.key",
			ClientAuth: tls.RequireAndVerifyClientCert,
		},
	}

	d := spawnDaemon(t, conf)
	defer d.Close()

	clientWithCert := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: conf.TLS.ClientTLS,
		},
	}

	clientWithoutCert := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				RootCAs: conf.TLS.ServerTLS.RootCAs,
			},
		},
	}

	reqCertRequired, err := http.NewRequest(http.MethodGet, fmt.Sprintf("https://%s/v1/HealthCheck", conf.HTTPListenAddress), nil)
	require.NoError(t, err)
	reqNoClientCertRequired, err := http.NewRequest(http.MethodGet, fmt.Sprintf("https://%s/v1/HealthCheck", conf.HTTPStatusListenAddress), nil)
	require.NoError(t, err)

	// Test that a client without a cert can access /v1/HealthCheck at status address
	resp, err := clientWithoutCert.Do(reqNoClientCertRequired)
	require.NoError(t, err)
	defer resp.Body.Close()
	b, err := io.ReadAll(resp.Body)
	require.NoError(t, err)
	assert.Equal(t, `{"status":"healthy","message":"","peer_count":1}`, strings.ReplaceAll(string(b), " ", ""))

	// Verify we get an error when we try to access existing HTTPListenAddress without cert
	//nolint:bodyclose // Expect error, no body to close.
	_, err = clientWithoutCert.Do(reqCertRequired)
	require.Error(t, err)

	// Check that with a valid client cert we can access /v1/HealthCheck at existing HTTPListenAddress
	resp2, err := clientWithCert.Do(reqCertRequired)
	require.NoError(t, err)
	defer resp2.Body.Close()
	b, err = io.ReadAll(resp2.Body)
	require.NoError(t, err)
	assert.Equal(t, `{"status":"healthy","message":"","peer_count":1}`, strings.ReplaceAll(string(b), " ", ""))
}