1
package server
2
 
3
import (
4
	"bytes"
5
	"context"
6
	"fmt"
7
	"net"
8
	"sync"
9
	"time"
10
 
11
	"go.miragespace.co/specter/spec/protocol"
12
	"go.miragespace.co/specter/spec/transport"
13
 
14
	"github.com/twitchtv/twirp"
15
	"github.com/zhangyunhao116/skipmap"
16
)
17
 
18
const (
19
	maxSessions           = 1024
20
	openFailureCloseDelay = time.Second
21
)
22
 
23
var (
24
	revalidateInterval = 30 * time.Second
25
	revalidateGrace    = 2 * time.Minute
26
)
27
 
28
type sessionMode uint8
29
 
30
const (
31
	ephemeral sessionMode = iota
32
	delegated
33
)
34
 
35
type sessionState uint8
36
 
37
const (
38
	provisional sessionState = iota
39
	active
40
	terminal
41
)
42
 
43
type session struct {
44
	alias            string
45
	hostname         string
46
	conn             transport.PhysicalConn
47
	mode             sessionMode
48
	spki             []byte
49
	grantID          string
50
	owner            *protocol.ClientToken
51
	grantExpiry      time.Time
52
	certNotAfter     time.Time
53
	mu               sync.Mutex
54
	state            sessionState
55
	authorizedUntil  time.Time
56
	revalidateCancel context.CancelFunc
57
	revalidateDone   chan struct{}
58
}
59
 
60
type sessionRegistry struct {
61
	// mu protects admission, byConn, and conditional alias updates.
62
	mu      sync.Mutex
63
	byAlias *skipmap.StringMap[*session]
64
	byConn  map[transport.PhysicalConn]*session
65
}
66
 
67
func newSessionRegistry() *sessionRegistry {
68
	return &sessionRegistry{
69
		byAlias: skipmap.NewString[*session](),
70
		byConn:  make(map[transport.PhysicalConn]*session),
71
	}
72
}
73
 
74
// describe joins operator observations to the exact attachment. Certificate
75
// identity alone cannot distinguish a stale connection from its replacement.
76
func (r *sessionRegistry) describe(conn transport.PhysicalConn) (mode, hostname, owner string) {
77
	if r == nil || conn == nil {
78
		return "", "", ""
79
	}
80
	r.mu.Lock()
81
	s := r.byConn[conn]
82
	r.mu.Unlock()
83
	if s == nil {
84
		return "", "", ""
85
	}
86
	s.mu.Lock()
87
	defer s.mu.Unlock()
88
	if s.state != active || connectionDone(s.conn) {
89
		return "", "", ""
90
	}
91
	if s.mode == ephemeral {
92
		return "ephemeral", s.hostname, ""
93
	}
94
	return "token", s.hostname, string(s.owner.GetToken())
95
}
96
 
97
func connectionDone(c transport.PhysicalConn) bool {
98
	select {
99
	case <-c.Done():
100
		return true
101
	default:
102
		return false
103
	}
104
}
105
 
106
func (r *sessionRegistry) reserve(s *session) error {
107
	r.mu.Lock()
108
	defer r.mu.Unlock()
109
	if r.byConn[s.conn] != nil {
110
		return twirp.FailedPrecondition.Error("one Open attempt is allowed per attachment")
111
	}
112
	if len(r.byConn) >= maxSessions {
113
		return twirp.ResourceExhausted.Error("lightweight session limit reached")
114
	}
115
	// Failed attempts also occupy this attachment until physical termination.
116
	r.byConn[s.conn] = s
117
	go r.watch(s)
118
	if old, ok := r.byAlias.Load(s.alias); ok {
119
		old.mu.Lock()
120
		defer old.mu.Unlock()
121
		if s.mode == ephemeral && (!bytes.Equal(old.spki, s.spki) || old.hostname != s.hostname) {
122
			return twirp.PermissionDenied.Error("ephemeral identity collision")
123
		}
124
		if old.state != terminal && !connectionDone(old.conn) {
125
			return twirp.Unavailable.Error("session alias is still attached")
126
		}
127
	}
128
	r.byAlias.Store(s.alias, s)
129
	return nil
130
}
131
 
132
func (r *sessionRegistry) watch(s *session) {
133
	<-s.conn.Done()
134
	s.mu.Lock()
135
	s.state = terminal
136
	cancel := s.revalidateCancel
137
	done := s.revalidateDone
138
	s.mu.Unlock()
139
	if cancel != nil {
140
		cancel()
141
		// Keep unfinished storage work in the admission count until it returns.
142
		<-done
143
	}
144
	r.mu.Lock()
145
	defer r.mu.Unlock()
146
	if current, ok := r.byAlias.Load(s.alias); ok && current == s {
147
		r.byAlias.Delete(s.alias)
148
	}
149
	if r.byConn[s.conn] == s {
150
		delete(r.byConn, s.conn)
151
	}
152
}
153
 
154
func (s *session) activate(ctx context.Context, until time.Time) bool {
155
	s.mu.Lock()
156
	defer s.mu.Unlock()
157
	if s.state != provisional || connectionDone(s.conn) || ctx.Err() != nil || (s.mode == delegated && !time.Now().Before(until)) {
158
		return false
159
	}
160
	s.state, s.authorizedUntil = active, until
161
	return true
162
}
163
 
164
func (s *session) fail(reason string) {
165
	s.mu.Lock()
166
	s.state = terminal
167
	s.mu.Unlock()
168
	// Allow the unary error response to flush before ending its connection.
169
	time.AfterFunc(openFailureCloseDelay, func() { s.conn.Close(reason) })
170
}
171
 
172
func (s *session) close(reason string) {
173
	s.mu.Lock()
174
	s.state = terminal
175
	s.mu.Unlock()
176
	s.conn.Close(reason)
177
}
178
 
179
func (r *sessionRegistry) dial(alias, hostname string) (net.Conn, error) {
180
	s, ok := r.byAlias.Load(alias)
181
	if !ok {
182
		return nil, fmt.Errorf("%w: unknown session", transport.ErrNoDirect)
183
	}
184
	s.mu.Lock()
185
	defer s.mu.Unlock()
186
	if s.state != active || s.hostname != hostname || connectionDone(s.conn) || (s.mode == delegated && !time.Now().Before(s.authorizedUntil)) {
187
		return nil, fmt.Errorf("%w: session is not authorized", transport.ErrNoDirect)
188
	}
189
	conn, err := s.conn.OpenStream(protocol.Stream_DIRECT)
190
	if err != nil {
191
		return nil, fmt.Errorf("%w: opening session stream: %v", transport.ErrNoDirect, err)
192
	}
193
	return conn, nil
194
}
195
 
196
func (s *session) authorizationDeadline(start, expiry time.Time) time.Time {
197
	until := start.Add(revalidateGrace)
198
	if !expiry.IsZero() && expiry.Before(until) {
199
		until = expiry
200
	}
201
	if s.certNotAfter.Before(until) {
202
		until = s.certNotAfter
203
	}
204
	// Preserve monotonic elapsed-time checks even for wall-clock expiry fields.
205
	return start.Add(until.Sub(start))
206
}
207
 
208
func (s *session) extend(start, expiry time.Time) {
209
	s.mu.Lock()
210
	defer s.mu.Unlock()
211
	until := s.authorizationDeadline(start, expiry)
212
	if s.state == active && time.Now().Before(s.authorizedUntil) && until.After(s.authorizedUntil) {
213
		s.authorizedUntil = until
214
	}
215
}