Skip to content

Commit

Permalink
Feature to override http client
Browse files Browse the repository at this point in the history
This feature could be used for multiple purposes:
from custom retry to ssl CA injection

Closes #29
  • Loading branch information
dkropachev authored and nyh committed Nov 18, 2024
1 parent c6f2c02 commit 6d02241
Show file tree
Hide file tree
Showing 2 changed files with 42 additions and 6 deletions.
19 changes: 15 additions & 4 deletions go/v2/alternator_lb.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@ import (
smithyhttp "github.com/aws/smithy-go/transport/http"
)

type httpClient interface {
Do(req *http.Request) (*http.Response, error)
}

// AlternatorNodes holds the configuration for the load balanced alternator nodes as well as some locks for the load balancing thread.
type AlternatorNodes struct {
fetchInterval time.Duration
Expand All @@ -26,6 +30,7 @@ type AlternatorNodes struct {

nodes []string
nextNodeIndex int
httpClient httpClient

mutex sync.Mutex
ctx context.Context
Expand All @@ -35,12 +40,18 @@ type AlternatorNodes struct {
// NewAlternatorNodes creates a new, unstarted instance of the alternator nodes loadbalancing.
func NewAlternatorNodes(scheme string, port int, initialNodes ...string) *AlternatorNodes {
return &AlternatorNodes{
scheme: scheme,
port: port,
nodes: initialNodes,
scheme: scheme,
port: port,
nodes: initialNodes,
httpClient: http.DefaultClient,
}
}

// SetHTTPClient changes underlying http client used for fetching the list of alternator nodes.
func (n *AlternatorNodes) SetHTTPClient(client httpClient) {
n.httpClient = client
}

// Config produces a conf for the AWS SDK that will integrate the alternator loadbalancing with the AWS SDK.
func (n *AlternatorNodes) Config(fake_domain string, key string, secret_key string) aws.Config {
fake_url := fmt.Sprintf("%s://%s:%d", n.scheme, fake_domain, n.port)
Expand Down Expand Up @@ -116,7 +127,7 @@ func (n *AlternatorNodes) doUpdateNodes() error {
return err
}

resp, err := http.DefaultClient.Do(req)
resp, err := n.httpClient.Do(req)
if err != nil {
return err
}
Expand Down
29 changes: 27 additions & 2 deletions go/v2/example_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,17 @@ package alternatorlb_test

import (
"context"
"crypto/tls"
"crypto/x509"
"fmt"
"time"

"github.com/aws/aws-sdk-go-v2/service/dynamodb"
alternatorlb "github.com/scylladb/alternator-load-balancing/go/v2"
"net/http"
"time"
)

var customPEMCertificate = []byte(nil)

func ExampleAlternatorNodes() {
ctx := context.Background()

Expand All @@ -22,6 +26,27 @@ func ExampleAlternatorNodes() {
// Use the local Alternator with our silly testing alternator/secret_pass
// authentication - and the new load balancing code.
alternatorNodes := alternatorlb.NewAlternatorNodes("http", 8000, "127.0.0.1")

// To add custom CA certificate on top of system CA certificates:
if customPEMCertificate != nil {
systemPool, err := x509.SystemCertPool()
if err != nil {
panic(err)
}

if !systemPool.AppendCertsFromPEM(customPEMCertificate) {
panic("failed to append custom certificate")
}

alternatorNodes.SetHTTPClient(&http.Client{
Transport: &http.Transport{
TLSClientConfig: &tls.Config{
RootCAs: systemPool,
},
},
})
}

alternatorNodes.Start(ctx, 1*time.Second)
defer alternatorNodes.Stop()

Expand Down

0 comments on commit 6d02241

Please sign in to comment.