Skip to content

Commit 0a48280

Browse files
h4x3rotabh4x3rotab
committed
Harden mesh startup sequencing
1 parent b79d7ed commit 0a48280

3 files changed

Lines changed: 137 additions & 14 deletions

File tree

‎consul-postgres-ha/admission-client/main.go‎

Lines changed: 57 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"crypto/sha512"
77
"encoding/hex"
88
"encoding/json"
9+
"errors"
910
"flag"
1011
"fmt"
1112
"io"
@@ -19,6 +20,8 @@ import (
1920
dstack "github.com/Dstack-TEE/dstack/sdk/go/dstack"
2021
)
2122

23+
var admissionRetryInterval = 2 * time.Second
24+
2225
func main() {
2326
var (
2427
identity = flag.String("identity", os.Getenv("ADMISSION_IDENTITY"), "claimed SPIFFE identity")
@@ -96,15 +99,30 @@ func admit(ctx context.Context, httpClient *http.Client, dstackClient quoteClien
9699
if len(brokerURLs) == 0 {
97100
return "", fmt.Errorf("no broker URLs configured")
98101
}
99-
var errs []string
100-
for _, brokerURL := range brokerURLs {
101-
token, err := admitOne(ctx, httpClient, dstackClient, strings.TrimRight(brokerURL, "/"), identity, statementBytes, vmConfig)
102-
if err == nil {
103-
return token, nil
102+
var lastErrs []string
103+
for {
104+
if err := ctx.Err(); err != nil {
105+
if len(lastErrs) > 0 {
106+
return "", fmt.Errorf("admission timed out or was canceled after broker errors: %s", strings.Join(lastErrs, "; "))
107+
}
108+
return "", err
109+
}
110+
lastErrs = lastErrs[:0]
111+
for _, brokerURL := range brokerURLs {
112+
brokerURL = strings.TrimRight(brokerURL, "/")
113+
token, err := admitOne(ctx, httpClient, dstackClient, brokerURL, identity, statementBytes, vmConfig)
114+
if err == nil {
115+
return token, nil
116+
}
117+
if isPermanentAdmissionError(err) {
118+
return "", fmt.Errorf("%s: %w", brokerURL, err)
119+
}
120+
lastErrs = append(lastErrs, fmt.Sprintf("%s: %v", brokerURL, err))
121+
}
122+
if !sleepContext(ctx, admissionRetryInterval) {
123+
return "", fmt.Errorf("admission timed out or was canceled after broker errors: %s", strings.Join(lastErrs, "; "))
104124
}
105-
errs = append(errs, fmt.Sprintf("%s: %v", brokerURL, err))
106125
}
107-
return "", fmt.Errorf("all admission brokers rejected or failed: %s", strings.Join(errs, "; "))
108126
}
109127

110128
func admitOne(ctx context.Context, httpClient *http.Client, dstackClient quoteClient, brokerURL, identity string, statementBytes []byte, vmConfig string) (string, error) {
@@ -141,7 +159,7 @@ func challenge(ctx context.Context, httpClient *http.Client, brokerURL string) (
141159
defer resp.Body.Close()
142160
if resp.StatusCode != http.StatusOK {
143161
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
144-
return "", fmt.Errorf("challenge HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
162+
return "", httpStatusError{status: resp.StatusCode, body: strings.TrimSpace(string(body))}
145163
}
146164
var out struct {
147165
Nonce string `json:"nonce"`
@@ -180,7 +198,7 @@ func attest(ctx context.Context, httpClient *http.Client, brokerURL string, payl
180198
defer resp.Body.Close()
181199
if resp.StatusCode != http.StatusOK {
182200
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
183-
return "", fmt.Errorf("attest HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
201+
return "", httpStatusError{status: resp.StatusCode, body: strings.TrimSpace(string(body))}
184202
}
185203
var out struct {
186204
ConsulACLToken string `json:"consul_acl_token"`
@@ -232,3 +250,33 @@ func shortHash(h string) string {
232250
}
233251
return h[:12]
234252
}
253+
254+
type httpStatusError struct {
255+
status int
256+
body string
257+
}
258+
259+
func (e httpStatusError) Error() string {
260+
return fmt.Sprintf("HTTP %d: %s", e.status, e.body)
261+
}
262+
263+
func isPermanentAdmissionError(err error) bool {
264+
var statusErr httpStatusError
265+
if !errors.As(err, &statusErr) {
266+
return false
267+
}
268+
return statusErr.status == http.StatusBadRequest ||
269+
statusErr.status == http.StatusUnauthorized ||
270+
statusErr.status == http.StatusForbidden
271+
}
272+
273+
func sleepContext(ctx context.Context, d time.Duration) bool {
274+
timer := time.NewTimer(d)
275+
defer timer.Stop()
276+
select {
277+
case <-timer.C:
278+
return true
279+
case <-ctx.Done():
280+
return false
281+
}
282+
}

‎consul-postgres-ha/admission-client/main_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,16 @@ package main
22

33
import (
44
"bytes"
5+
"context"
56
"crypto/sha512"
67
"encoding/hex"
8+
"net/http"
9+
"net/http/httptest"
10+
"sync/atomic"
711
"testing"
12+
"time"
13+
14+
dstack "github.com/Dstack-TEE/dstack/sdk/go/dstack"
815
)
916

1017
func TestReportDataBindsStatementAndNonce(t *testing.T) {
@@ -36,3 +43,68 @@ func TestSplitCSV(t *testing.T) {
3643
t.Fatalf("bad split: %#v", got)
3744
}
3845
}
46+
47+
func TestAdmitRetriesTransientBrokerFailure(t *testing.T) {
48+
prevRetry := admissionRetryInterval
49+
admissionRetryInterval = 10 * time.Millisecond
50+
t.Cleanup(func() { admissionRetryInterval = prevRetry })
51+
52+
var challenges atomic.Int32
53+
broker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
54+
switch r.URL.Path {
55+
case "/v1/admission/challenge":
56+
if challenges.Add(1) == 1 {
57+
http.Error(w, "not ready", http.StatusBadGateway)
58+
return
59+
}
60+
w.Write([]byte(`{"nonce":"000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f"}`))
61+
case "/v1/admission/attest":
62+
w.Write([]byte(`{"consul_acl_token":"issued-token"}`))
63+
default:
64+
http.NotFound(w, r)
65+
}
66+
}))
67+
t.Cleanup(broker.Close)
68+
69+
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
70+
defer cancel()
71+
token, err := admit(ctx, broker.Client(), fakeQuoteClient{}, []string{broker.URL}, "spiffe://demo/webdemo", []byte(`{"ok":true}`), "{}")
72+
if err != nil {
73+
t.Fatal(err)
74+
}
75+
if token != "issued-token" {
76+
t.Fatalf("wrong token: %q", token)
77+
}
78+
if got := challenges.Load(); got != 2 {
79+
t.Fatalf("expected one retry, got %d challenge requests", got)
80+
}
81+
}
82+
83+
func TestAdmitDoesNotRetryPermanentBrokerRejection(t *testing.T) {
84+
prevRetry := admissionRetryInterval
85+
admissionRetryInterval = 10 * time.Millisecond
86+
t.Cleanup(func() { admissionRetryInterval = prevRetry })
87+
88+
var challenges atomic.Int32
89+
broker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
90+
challenges.Add(1)
91+
http.Error(w, "admission rejected", http.StatusForbidden)
92+
}))
93+
t.Cleanup(broker.Close)
94+
95+
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
96+
defer cancel()
97+
_, err := admit(ctx, broker.Client(), fakeQuoteClient{}, []string{broker.URL}, "spiffe://demo/webdemo", []byte(`{"ok":true}`), "{}")
98+
if err == nil {
99+
t.Fatal("expected error")
100+
}
101+
if got := challenges.Load(); got != 1 {
102+
t.Fatalf("permanent rejection should not be retried, got %d challenge requests", got)
103+
}
104+
}
105+
106+
type fakeQuoteClient struct{}
107+
108+
func (fakeQuoteClient) Attest(context.Context, []byte) (*dstack.AttestResponse, error) {
109+
return &dstack.AttestResponse{Attestation: []byte("attestation")}, nil
110+
}

‎consul-postgres-ha/mesh-sidecar/entrypoint.sh‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -204,11 +204,13 @@ if [ "$ROLE" = "coordinator" ] && [ "${ADMISSION_BROKER_ENABLE:-}" = "1" ]; then
204204
ADMISSION_BROKER_PID=$!
205205
fi
206206

207-
# Wait for the local agent to accept HTTP requests; everything below
208-
# (sidecar registration, leaf cert, envoy bootstrap) goes through it.
207+
# Wait for the local agent to listen on its HTTP socket. This is deliberately
208+
# a transport-level local check, not a Consul API authorization check and not
209+
# cluster membership: workers cannot join servers until mesh-conn is up, and
210+
# ACL-protected agent endpoints can return primary-dc-down before that join.
209211
wait_consul_ready() {
210212
local n=0
211-
until CONSUL_HTTP_TOKEN="${CONSUL_MANAGEMENT_TOKEN:-}" consul members >/dev/null 2>&1; do
213+
until timeout 1 bash -c "cat < /dev/null > /dev/tcp/${CONSUL_HTTP%:*}/${CONSUL_HTTP##*:}" 2>/dev/null; do
212214
n=$((n+1))
213215
if [ $n -gt 60 ]; then
214216
log "consul agent not reachable after 60s"
@@ -283,7 +285,8 @@ if [ "$ROLE" = "worker" ]; then
283285
-broker-urls "$BROKER_URLS" \
284286
-token-file "$TOKEN_FILE" \
285287
-cluster "$CLUSTER_NAME" \
286-
-peer-id "$PEER_ID" 2>&1 | prefix "admission-client-${PARENT}"
288+
-peer-id "$PEER_ID" \
289+
-timeout 15m 2>&1 | prefix "admission-client-${PARENT}"
287290
BACKEND_CONSUL_TOKEN=$(cat "$TOKEN_FILE")
288291
fi
289292

@@ -424,7 +427,7 @@ if [ "$ROLE" = "coordinator" ] && [ "$ORDINAL" = "0" ]; then
424427
wait_consul_ready
425428
# Wait for the full coordinator quorum to be visible — config-entry
426429
# writes need a server with leadership and a healthy quorum.
427-
until [ "$(consul members 2>/dev/null | awk 'NR>1 && $4=="server" && $3=="alive"' | wc -l)" -ge "${BOOTSTRAP_EXPECT}" ]; do
430+
until [ "$(CONSUL_HTTP_TOKEN="${CONSUL_MANAGEMENT_TOKEN:-}" consul members 2>/dev/null | awk 'NR>1 && $4=="server" && $3=="alive"' | wc -l)" -ge "${BOOTSTRAP_EXPECT}" ]; do
428431
sleep 2
429432
done
430433
log "config-entry writer: quorum ready, writing entries"

0 commit comments

Comments
 (0)