282 lines
8.2 KiB
Go
282 lines
8.2 KiB
Go
package localreset
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
type testSession struct {
|
|
selected atomic.Int32
|
|
verified atomic.Int32
|
|
resets atomic.Int32
|
|
closed atomic.Int32
|
|
}
|
|
|
|
func (s *testSession) Methods() []Method { return []Method{{ID: "recovery", Label: "恢复码"}} }
|
|
func (s *testSession) Select(ctx context.Context, method string) error {
|
|
if ctx.Err() != nil || method != "recovery" {
|
|
return ErrProtocol
|
|
}
|
|
s.selected.Add(1)
|
|
return nil
|
|
}
|
|
func (s *testSession) Verify(ctx context.Context, code string) error {
|
|
if ctx.Err() != nil || s.selected.Load() != 1 || code != "test-code" {
|
|
return ErrProtocol
|
|
}
|
|
s.verified.Add(1)
|
|
return nil
|
|
}
|
|
func (s *testSession) Reset(ctx context.Context, password string) error {
|
|
if ctx.Err() != nil || s.verified.Load() != 1 || password != "Example123!" {
|
|
return ErrProtocol
|
|
}
|
|
s.resets.Add(1)
|
|
return nil
|
|
}
|
|
func (s *testSession) Close() { s.closed.Add(1) }
|
|
|
|
func testConnection(t *testing.T, timeout time.Duration) (*Server, net.Conn, *testSession, <-chan struct{}) {
|
|
t.Helper()
|
|
serverConn, clientConn := net.Pipe()
|
|
server := newServer(nil, func(net.Conn) (Peer, error) { return Peer{Identity: "test"}, nil }, nil)
|
|
server.timeout = timeout
|
|
session := &testSession{}
|
|
done := make(chan struct{})
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
go func() {
|
|
defer close(done)
|
|
defer serverConn.Close()
|
|
server.handle(ctx, serverConn, func(context.Context, Peer) (Session, error) { return session, nil })
|
|
}()
|
|
t.Cleanup(func() {
|
|
cancel()
|
|
_ = clientConn.Close()
|
|
_ = serverConn.Close()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Error("connection did not stop")
|
|
}
|
|
})
|
|
return server, clientConn, session, done
|
|
}
|
|
|
|
func TestProtocolRoundTrip(t *testing.T) {
|
|
_, conn, session, done := testConnection(t, time.Second)
|
|
client, err := Begin(context.Background(), conn)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer client.Close()
|
|
methods := client.Methods()
|
|
methods[0].ID = "mutated"
|
|
if client.Methods()[0].ID != "recovery" {
|
|
t.Fatal("methods leaked mutable state")
|
|
}
|
|
if err := client.Select("recovery"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := client.Verify("test-code"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := client.Reset("Example123!"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := client.Reset("Example123!"); err == nil {
|
|
t.Fatal("replay accepted")
|
|
}
|
|
<-done
|
|
if session.resets.Load() != 1 || session.closed.Load() != 1 {
|
|
t.Fatal("session not consumed and closed once")
|
|
}
|
|
}
|
|
|
|
func TestProtocolRejectsInvalidFrames(t *testing.T) {
|
|
cases := map[string]func(string) string{
|
|
"invalid-json": func(string) string { return "{broken}\n" },
|
|
"oversized": func(string) string { return strings.Repeat("x", MaxFrameSize+1) + "\n" },
|
|
"wrong-version": func(n string) string {
|
|
return `{"version":9,"nonce":"` + n + `","action":"select","method":"recovery"}` + "\n"
|
|
},
|
|
"wrong-nonce": func(string) string {
|
|
return `{"version":1,"nonce":"wrong","action":"select","method":"recovery"}` + "\n"
|
|
},
|
|
"skip-verification": func(n string) string {
|
|
return `{"version":1,"nonce":"` + n + `","action":"reset","password":"Example123!"}` + "\n"
|
|
},
|
|
"unknown-field": func(n string) string {
|
|
return `{"version":1,"nonce":"` + n + `","action":"select","method":"recovery","extra":true}` + "\n"
|
|
},
|
|
"trailing-object": func(n string) string {
|
|
return `{"version":1,"nonce":"` + n + `","action":"select","method":"recovery"}{}` + "\n"
|
|
},
|
|
"wrong-method": func(n string) string {
|
|
return `{"version":1,"nonce":"` + n + `","action":"select","method":"missing"}` + "\n"
|
|
},
|
|
"mixed-payload": func(n string) string {
|
|
return `{"version":1,"nonce":"` + n + `","action":"select","method":"recovery","password":"Example123!"}` + "\n"
|
|
},
|
|
}
|
|
for name, payload := range cases {
|
|
t.Run(name, func(t *testing.T) {
|
|
_, conn, session, done := testConnection(t, time.Second)
|
|
reader := bufio.NewReaderSize(conn, MaxFrameSize+1)
|
|
var challenge Response
|
|
if err := readFrame(reader, &challenge); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
go func() { _, _ = io.WriteString(conn, payload(challenge.Nonce)) }()
|
|
var response Response
|
|
if err := readFrame(reader, &response); err != nil || response.OK {
|
|
t.Fatalf("accepted malformed request: %+v %v", response, err)
|
|
}
|
|
<-done
|
|
if session.resets.Load() != 0 || session.closed.Load() != 1 {
|
|
t.Fatal("invalid request reached reset or leaked session")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestProtocolRotatesNonce(t *testing.T) {
|
|
_, conn, session, done := testConnection(t, time.Second)
|
|
reader := bufio.NewReaderSize(conn, MaxFrameSize+1)
|
|
var first, second, failure Response
|
|
if err := readFrame(reader, &first); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := writeFrame(conn, Request{Version: Version, Nonce: first.Nonce, Action: "select", Method: "recovery"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := readFrame(reader, &second); err != nil || second.Nonce == first.Nonce {
|
|
t.Fatal("nonce not rotated")
|
|
}
|
|
if err := writeFrame(conn, Request{Version: Version, Nonce: first.Nonce, Action: "verify", Code: "test-code"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := readFrame(reader, &failure); err != nil || failure.OK {
|
|
t.Fatal("replayed nonce accepted")
|
|
}
|
|
<-done
|
|
if session.verified.Load() != 0 {
|
|
t.Fatal("replay reached verification")
|
|
}
|
|
}
|
|
|
|
func TestProtocolDeadlineClosesIdleConnection(t *testing.T) {
|
|
_, conn, session, done := testConnection(t, 30*time.Millisecond)
|
|
var response Response
|
|
if err := readFrame(bufio.NewReaderSize(conn, MaxFrameSize+1), &response); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("idle connection survived deadline")
|
|
}
|
|
if session.closed.Load() != 1 {
|
|
t.Fatal("expired session leaked")
|
|
}
|
|
}
|
|
|
|
func TestProtocolVerificationFailureIsTerminal(t *testing.T) {
|
|
_, conn, session, done := testConnection(t, time.Second)
|
|
client, err := Begin(context.Background(), conn)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := client.Select("recovery"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := client.Verify("wrong-code"); err == nil {
|
|
t.Fatal("wrong code accepted")
|
|
}
|
|
if err := client.Reset("Example123!"); err == nil {
|
|
t.Fatal("failed session reused")
|
|
}
|
|
<-done
|
|
if session.resets.Load() != 0 {
|
|
t.Fatal("verification bypassed")
|
|
}
|
|
}
|
|
|
|
func TestFailureWindowExpires(t *testing.T) {
|
|
server := newServer(nil, nil, nil)
|
|
now := time.Now()
|
|
server.now = func() time.Time { return now }
|
|
for i := 0; i < 5; i++ {
|
|
if !server.permit("peer") {
|
|
t.Fatal("premature limit")
|
|
}
|
|
server.failed("peer")
|
|
}
|
|
if server.permit("peer") {
|
|
t.Fatal("failure limit bypassed")
|
|
}
|
|
now = now.Add(time.Minute)
|
|
if !server.permit("peer") {
|
|
t.Fatal("failure limit did not expire")
|
|
}
|
|
}
|
|
|
|
func TestClientCancellationClosesConnection(t *testing.T) {
|
|
_, conn, session, done := testConnection(t, time.Minute)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
client, err := Begin(ctx, conn)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer client.Close()
|
|
cancel()
|
|
select {
|
|
case <-done:
|
|
case <-time.After(time.Second):
|
|
t.Fatal("cancellation did not close connection")
|
|
}
|
|
if session.closed.Load() != 1 {
|
|
t.Fatal("cancelled session leaked")
|
|
}
|
|
}
|
|
|
|
func TestFactoryErrorDoesNotDiscloseDetails(t *testing.T) {
|
|
left, right := net.Pipe()
|
|
defer right.Close()
|
|
server := newServer(nil, func(net.Conn) (Peer, error) { return Peer{Identity: "test"}, nil }, nil)
|
|
go func() {
|
|
defer left.Close()
|
|
server.handle(context.Background(), left, func(context.Context, Peer) (Session, error) {
|
|
return nil, errors.New("secret detail")
|
|
})
|
|
}()
|
|
data, err := io.ReadAll(right)
|
|
if err != nil || strings.Contains(string(data), "secret") || !strings.Contains(string(data), "UNAVAILABLE") {
|
|
t.Fatalf("unsafe response: %s %v", data, err)
|
|
}
|
|
}
|
|
|
|
func TestUnauthorizedPeerCannotCreateSession(t *testing.T) {
|
|
left, right := net.Pipe()
|
|
defer right.Close()
|
|
server := newServer(nil, func(net.Conn) (Peer, error) { return Peer{}, ErrUnauthorized }, nil)
|
|
var calls atomic.Int32
|
|
go func() {
|
|
defer left.Close()
|
|
server.handle(context.Background(), left, func(context.Context, Peer) (Session, error) {
|
|
calls.Add(1)
|
|
return &testSession{}, nil
|
|
})
|
|
}()
|
|
data, err := io.ReadAll(right)
|
|
if err != nil || !strings.Contains(string(data), "UNAUTHORIZED") || calls.Load() != 0 {
|
|
t.Fatal("unauthorized peer reached factory")
|
|
}
|
|
}
|