cdrop — 跨 OS 剪贴板与文件传输服务
Commilitia Drop：自托管的跨设备剪贴板同步与点对点文件传输。 - 后端 Go（chi / SQLite WAL / SSE Hub / WebRTC signaling + 状态机 / Relay ring buffer），编译进单个 distroless 镜像（前端 go:embed）。 - 前端 React + TanStack Router + Zustand，自实现 SSE + WebRTC P2P，NAT 受阻时回退服务端中继；聚珍（Juzhen）CJK 综合排版。 - 桌面端 Wails v2（macOS / Windows），瘦客户端复用 web。 - 鉴权 OIDC PKCE（自建 Casdoor 等），refresh_token 信封加密存系统密钥库；iOS Shortcut 用 HS256 scoped token。 架构文档与变更记录见 docs 分支（PROJECT_BRIEF / FRONTEND_DESIGN / CHANGELOG）。 本次为公开发布初始提交：完整开发历史（含部署细节）留存于私有归档，公开仓库自此干净起步。
This commit is contained in:
@@ -0,0 +1,269 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultMaxSessions = 8
|
||||
DefaultMaxSessionsPerUser = 3
|
||||
DefaultSessionCap = 64 << 20 // 64 MiB
|
||||
DefaultChunkLimit = 1 << 20 // 1 MiB
|
||||
|
||||
// DefaultIdleTTL is how long a registered relay session may sit without any
|
||||
// chunk/stream activity before the reaper tears it down. A fallback that
|
||||
// registered but never got used (both peers vanished) is reclaimed rather
|
||||
// than pinning a slot — and, once bytes flowed, its 64 MiB ring — forever.
|
||||
DefaultIdleTTL = 2 * time.Minute
|
||||
|
||||
// DefaultReapPeriod is how often the reaper goroutine sweeps for idle sessions.
|
||||
DefaultReapPeriod = 30 * time.Second
|
||||
)
|
||||
|
||||
var (
|
||||
ErrTooManySessions = errors.New("relay: too many active sessions")
|
||||
ErrForbidden = errors.New("relay: caller is not a participant of this session")
|
||||
ErrNotRegistered = errors.New("relay: session not registered")
|
||||
)
|
||||
|
||||
// Session pairs a Ring with the metadata that proves who is allowed to use it.
|
||||
// The ring is allocated lazily on the first Acquire so a registered-but-unused
|
||||
// fallback costs only bookkeeping until bytes actually flow.
|
||||
type Session struct {
|
||||
ID string
|
||||
UserID string
|
||||
Sender string // device name of the legitimate sender (POSTs chunks)
|
||||
Receiver string // device name of the legitimate receiver (GETs the stream)
|
||||
|
||||
capBytes int
|
||||
ring *Ring // nil until first Acquire; guarded by Manager.mu
|
||||
|
||||
bytesWritten atomic.Int64
|
||||
lastActivity atomic.Int64 // unix seconds; the reaper's idle clock
|
||||
readerActive atomic.Bool // single-consumer guard (a second reader tears the ring)
|
||||
}
|
||||
|
||||
func (s *Session) touch() { s.lastActivity.Store(time.Now().Unix()) }
|
||||
|
||||
func (s *Session) Write(ctx context.Context, p []byte) (int, error) {
|
||||
s.touch()
|
||||
n, err := s.ring.Write(ctx, p)
|
||||
if n > 0 {
|
||||
s.bytesWritten.Add(int64(n))
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (s *Session) Read(ctx context.Context, p []byte) (int, error) {
|
||||
s.touch()
|
||||
return s.ring.Read(ctx, p)
|
||||
}
|
||||
|
||||
func (s *Session) CloseWriter() {
|
||||
if s.ring != nil {
|
||||
s.ring.CloseWriter()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) Abort() {
|
||||
if s.ring != nil {
|
||||
s.ring.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) BytesWritten() int64 { return s.bytesWritten.Load() }
|
||||
|
||||
// AcquireReader claims the single consumer slot. The relay stream is strictly
|
||||
// single-consumer: two concurrent readers would each drain bytes out of the
|
||||
// ring and silently corrupt the receiver's copy (G2). A second reader gets a
|
||||
// false here so the handler can 409 instead of joining.
|
||||
func (s *Session) AcquireReader() bool { return s.readerActive.CompareAndSwap(false, true) }
|
||||
|
||||
// ReleaseReader frees the consumer slot when the stream handler returns.
|
||||
func (s *Session) ReleaseReader() { s.readerActive.Store(false) }
|
||||
|
||||
// Manager owns active relay sessions. A session must be Register-ed (which the
|
||||
// transfer service does when a transfer enters RELAY_ACTIVE) before either peer
|
||||
// can Acquire it. This is the core of the R1/G2 hardening: the old Acquire
|
||||
// minted a session for any unknown id on demand, so 8 one-byte requests could
|
||||
// exhaust every slot and pin ~512 MiB. Now only legitimate, server-minted
|
||||
// transfers that actually reached RELAY_ACTIVE get a slot, capped both globally
|
||||
// and per-user, and reclaimed when idle.
|
||||
type Manager struct {
|
||||
mu sync.Mutex
|
||||
sessions map[string]*Session
|
||||
perUser map[string]int
|
||||
maxSessions int
|
||||
maxPerUser int
|
||||
sessionCap int
|
||||
idleTTL time.Duration
|
||||
}
|
||||
|
||||
func NewManager(maxSessions, sessionCapBytes int) *Manager {
|
||||
if maxSessions <= 0 {
|
||||
maxSessions = DefaultMaxSessions
|
||||
}
|
||||
if sessionCapBytes <= 0 {
|
||||
sessionCapBytes = DefaultSessionCap
|
||||
}
|
||||
maxPerUser := DefaultMaxSessionsPerUser
|
||||
if maxPerUser > maxSessions {
|
||||
// A per-user cap above the global cap is meaningless; clamp so small
|
||||
// global caps (tests, constrained deploys) still behave sanely.
|
||||
maxPerUser = maxSessions
|
||||
}
|
||||
return &Manager{
|
||||
sessions: map[string]*Session{},
|
||||
perUser: map[string]int{},
|
||||
maxSessions: maxSessions,
|
||||
maxPerUser: maxPerUser,
|
||||
sessionCap: sessionCapBytes,
|
||||
idleTTL: DefaultIdleTTL,
|
||||
}
|
||||
}
|
||||
|
||||
// Register reserves a relay slot for a transfer that has just entered
|
||||
// RELAY_ACTIVE, recording its two legitimate participants. Only a registered
|
||||
// session can later be joined via Acquire.
|
||||
//
|
||||
// Idempotent for the same (id, userID): a duplicate Register just refreshes the
|
||||
// idle clock. A mismatched userID is ErrForbidden (ids are server-minted, so
|
||||
// this is defensive). Returns ErrTooManySessions when the global or per-user
|
||||
// cap is reached — the caller (transfer service) surfaces that as a retryable
|
||||
// failure rather than moving the transfer into a relay mode with no slot.
|
||||
func (m *Manager) Register(id, userID, sender, receiver string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if s, ok := m.sessions[id]; ok {
|
||||
if s.UserID != userID {
|
||||
return ErrForbidden
|
||||
}
|
||||
s.touch()
|
||||
return nil
|
||||
}
|
||||
if len(m.sessions) >= m.maxSessions {
|
||||
return ErrTooManySessions
|
||||
}
|
||||
if m.perUser[userID] >= m.maxPerUser {
|
||||
return ErrTooManySessions
|
||||
}
|
||||
|
||||
s := &Session{
|
||||
ID: id,
|
||||
UserID: userID,
|
||||
Sender: sender,
|
||||
Receiver: receiver,
|
||||
capBytes: m.sessionCap,
|
||||
}
|
||||
s.touch()
|
||||
m.sessions[id] = s
|
||||
m.perUser[userID] = m.perUser[userID] + 1
|
||||
return nil
|
||||
}
|
||||
|
||||
// Acquire joins an already-registered relay session as one of its two
|
||||
// participants. Unlike the old implementation it never creates a session:
|
||||
// - unknown id → ErrNotRegistered
|
||||
// - different user → ErrForbidden
|
||||
// - not sender/receiver → ErrForbidden
|
||||
//
|
||||
// device is the caller's X-Device-Name (set by the auth middleware). The 64 MiB
|
||||
// ring is allocated here, lazily, on the first join.
|
||||
func (m *Manager) Acquire(id, userID, device string) (*Session, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
s, ok := m.sessions[id]
|
||||
if !ok {
|
||||
return nil, ErrNotRegistered
|
||||
}
|
||||
if s.UserID != userID {
|
||||
return nil, ErrForbidden
|
||||
}
|
||||
if device == "" || (device != s.Sender && device != s.Receiver) {
|
||||
return nil, ErrForbidden
|
||||
}
|
||||
if s.ring == nil {
|
||||
s.ring = newRing(s.capBytes)
|
||||
}
|
||||
s.touch()
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// Release removes the session from the active set and aborts any blocked
|
||||
// Read/Write. Safe to call multiple times (terminal transitions and the pending
|
||||
// sweeper both call it). A no-op for an id that was never registered.
|
||||
func (m *Manager) Release(id string) {
|
||||
m.mu.Lock()
|
||||
s, ok := m.sessions[id]
|
||||
if ok {
|
||||
delete(m.sessions, id)
|
||||
m.decUserLocked(s.UserID)
|
||||
}
|
||||
m.mu.Unlock()
|
||||
if ok {
|
||||
s.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
// decUserLocked drops one from a user's active count, deleting the key at zero
|
||||
// so perUser never accumulates empty entries. Caller holds m.mu.
|
||||
func (m *Manager) decUserLocked(userID string) {
|
||||
if n := m.perUser[userID]; n <= 1 {
|
||||
delete(m.perUser, userID)
|
||||
} else {
|
||||
m.perUser[userID] = n - 1
|
||||
}
|
||||
}
|
||||
|
||||
// Active reports the current count of held sessions.
|
||||
func (m *Manager) Active() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return len(m.sessions)
|
||||
}
|
||||
|
||||
// reapIdle aborts and removes every session whose last activity is older than
|
||||
// the idle TTL, returning how many it reclaimed. now is unix seconds (injected
|
||||
// so tests don't depend on the wall clock).
|
||||
func (m *Manager) reapIdle(now int64) int {
|
||||
cutoff := now - int64(m.idleTTL/time.Second)
|
||||
var dead []*Session
|
||||
m.mu.Lock()
|
||||
for id, s := range m.sessions {
|
||||
if s.lastActivity.Load() < cutoff {
|
||||
delete(m.sessions, id)
|
||||
m.decUserLocked(s.UserID)
|
||||
dead = append(dead, s)
|
||||
}
|
||||
}
|
||||
m.mu.Unlock()
|
||||
for _, s := range dead {
|
||||
s.Abort()
|
||||
}
|
||||
return len(dead)
|
||||
}
|
||||
|
||||
// RunReaper blocks until ctx is cancelled, periodically reclaiming idle relay
|
||||
// sessions (fallbacks that registered but never completed — e.g. both peers
|
||||
// went away mid-transfer). Mirrors transfer.RunSweeper; start it once at boot.
|
||||
func (m *Manager) RunReaper(ctx context.Context) {
|
||||
t := time.NewTicker(DefaultReapPeriod)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
if n := m.reapIdle(time.Now().Unix()); n > 0 {
|
||||
slog.Info("relay reaper: reclaimed idle sessions", "count", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// register is a test helper: every session must be Register-ed before it can be
|
||||
// Acquire-d. Sender "s" / receiver "r" unless a test needs specific names.
|
||||
func register(t *testing.T, m *Manager, id, user string) {
|
||||
t.Helper()
|
||||
if err := m.Register(id, user, "s", "r"); err != nil {
|
||||
t.Fatalf("register %s/%s: %v", id, user, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_AcquireReturnsSameSession(t *testing.T) {
|
||||
m := NewManager(8, 1024)
|
||||
register(t, m, "s1", "alice")
|
||||
s1, err := m.Acquire("s1", "alice", "s")
|
||||
if err != nil {
|
||||
t.Fatalf("acquire 1: %v", err)
|
||||
}
|
||||
s2, err := m.Acquire("s1", "alice", "r")
|
||||
if err != nil {
|
||||
t.Fatalf("acquire 2: %v", err)
|
||||
}
|
||||
if s1 != s2 {
|
||||
t.Error("expected same session pointer on repeat Acquire")
|
||||
}
|
||||
}
|
||||
|
||||
// Acquire never mints a session: an id that was never Register-ed is rejected.
|
||||
// This is the core R1 fix — the old Acquire created one on demand, so any
|
||||
// authenticated client could conjure sessions for arbitrary ids and exhaust
|
||||
// every slot.
|
||||
func TestManager_AcquireUnregisteredRejected(t *testing.T) {
|
||||
m := NewManager(8, 1024)
|
||||
if _, err := m.Acquire("ghost", "alice", "s"); !errors.Is(err, ErrNotRegistered) {
|
||||
t.Errorf("expected ErrNotRegistered, got %v", err)
|
||||
}
|
||||
if got := m.Active(); got != 0 {
|
||||
t.Errorf("a rejected Acquire must not create a session; Active=%d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_DifferentUserForbidden(t *testing.T) {
|
||||
m := NewManager(8, 1024)
|
||||
register(t, m, "s1", "alice")
|
||||
if _, err := m.Acquire("s1", "mallory", "s"); !errors.Is(err, ErrForbidden) {
|
||||
t.Errorf("expected ErrForbidden, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Only the registered sender/receiver device may join (G2). A third device of
|
||||
// the same account — or a missing device name — is forbidden.
|
||||
func TestManager_NonParticipantForbidden(t *testing.T) {
|
||||
m := NewManager(8, 1024)
|
||||
if err := m.Register("s1", "alice", "laptop", "phone"); err != nil {
|
||||
t.Fatalf("register: %v", err)
|
||||
}
|
||||
if _, err := m.Acquire("s1", "alice", "tablet"); !errors.Is(err, ErrForbidden) {
|
||||
t.Errorf("non-participant device: expected ErrForbidden, got %v", err)
|
||||
}
|
||||
if _, err := m.Acquire("s1", "alice", ""); !errors.Is(err, ErrForbidden) {
|
||||
t.Errorf("empty device: expected ErrForbidden, got %v", err)
|
||||
}
|
||||
if _, err := m.Acquire("s1", "alice", "laptop"); err != nil {
|
||||
t.Errorf("sender device must be allowed: %v", err)
|
||||
}
|
||||
if _, err := m.Acquire("s1", "alice", "phone"); err != nil {
|
||||
t.Errorf("receiver device must be allowed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Register, not Acquire, now enforces the global cap.
|
||||
func TestManager_GlobalCapEnforced(t *testing.T) {
|
||||
m := NewManager(2, 1024)
|
||||
if err := m.Register("a", "u1", "s", "r"); err != nil {
|
||||
t.Fatalf("a: %v", err)
|
||||
}
|
||||
if err := m.Register("b", "u2", "s", "r"); err != nil {
|
||||
t.Fatalf("b: %v", err)
|
||||
}
|
||||
if err := m.Register("c", "u3", "s", "r"); !errors.Is(err, ErrTooManySessions) {
|
||||
t.Errorf("expected ErrTooManySessions, got %v", err)
|
||||
}
|
||||
if got := m.Active(); got != 2 {
|
||||
t.Errorf("Active: got %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A single user can't monopolise every global slot — the per-user cap stops
|
||||
// them well before the global cap, leaving room for other users.
|
||||
func TestManager_PerUserCapEnforced(t *testing.T) {
|
||||
m := NewManager(8, 1024) // per-user default is 3
|
||||
for i, id := range []string{"a", "b", "c"} {
|
||||
if err := m.Register(id, "greedy", "s", "r"); err != nil {
|
||||
t.Fatalf("register %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
if err := m.Register("d", "greedy", "s", "r"); !errors.Is(err, ErrTooManySessions) {
|
||||
t.Errorf("4th session for one user: expected ErrTooManySessions, got %v", err)
|
||||
}
|
||||
// A different user is unaffected — global slots remain.
|
||||
if err := m.Register("e", "other", "s", "r"); err != nil {
|
||||
t.Errorf("other user should still register: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Register is idempotent for the same owner (e.g. a retried /fallback) and does
|
||||
// not consume a second slot.
|
||||
func TestManager_RegisterIdempotent(t *testing.T) {
|
||||
m := NewManager(8, 1024)
|
||||
register(t, m, "s1", "alice")
|
||||
register(t, m, "s1", "alice")
|
||||
if got := m.Active(); got != 1 {
|
||||
t.Errorf("duplicate Register must not add a slot; Active=%d", got)
|
||||
}
|
||||
if err := m.Register("s1", "mallory", "s", "r"); !errors.Is(err, ErrForbidden) {
|
||||
t.Errorf("re-register under another user: expected ErrForbidden, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_ReleaseFreesSlot(t *testing.T) {
|
||||
m := NewManager(1, 1024)
|
||||
register(t, m, "a", "u")
|
||||
m.Release("a")
|
||||
if got := m.Active(); got != 0 {
|
||||
t.Errorf("Active after release: got %d, want 0", got)
|
||||
}
|
||||
// A new session can now slot in (both global and per-user counts dropped).
|
||||
if err := m.Register("b", "u", "s", "r"); err != nil {
|
||||
t.Errorf("register after release: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManager_ReleaseAbortsSession(t *testing.T) {
|
||||
m := NewManager(1, 1024)
|
||||
register(t, m, "a", "u")
|
||||
s, err := m.Acquire("a", "u", "s") // allocates the ring
|
||||
if err != nil {
|
||||
t.Fatalf("acquire: %v", err)
|
||||
}
|
||||
m.Release("a")
|
||||
if _, err := s.Write(nil, []byte("x")); !errors.Is(err, ErrClosed) {
|
||||
t.Errorf("write after release: got %v, want ErrClosed", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The reaper reclaims a session whose last activity is older than the idle TTL,
|
||||
// dropping both the global and per-user counts so the slot is reusable.
|
||||
func TestManager_ReaperReclaimsIdle(t *testing.T) {
|
||||
m := NewManager(8, 1024)
|
||||
register(t, m, "a", "u")
|
||||
|
||||
// Just-registered: a sweep "now" leaves it alone.
|
||||
if n := m.reapIdle(time.Now().Unix()); n != 0 {
|
||||
t.Errorf("fresh session reclaimed: got %d, want 0", n)
|
||||
}
|
||||
// Sweep far enough in the future that it's past the idle TTL.
|
||||
future := time.Now().Add(DefaultIdleTTL + time.Minute).Unix()
|
||||
if n := m.reapIdle(future); n != 1 {
|
||||
t.Errorf("idle session not reclaimed: got %d, want 1", n)
|
||||
}
|
||||
if got := m.Active(); got != 0 {
|
||||
t.Errorf("Active after reap: got %d, want 0", got)
|
||||
}
|
||||
// The per-user slot was freed too.
|
||||
if err := m.Register("b", "u", "s", "r"); err != nil {
|
||||
t.Errorf("register after reap: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The stream is single-consumer: a second concurrent reader is refused so it
|
||||
// can't drain bytes out from under the first and corrupt the receiver's copy.
|
||||
func TestSession_SingleReader(t *testing.T) {
|
||||
m := NewManager(8, 1024)
|
||||
register(t, m, "a", "u")
|
||||
s, err := m.Acquire("a", "u", "r")
|
||||
if err != nil {
|
||||
t.Fatalf("acquire: %v", err)
|
||||
}
|
||||
if !s.AcquireReader() {
|
||||
t.Fatal("first AcquireReader should succeed")
|
||||
}
|
||||
if s.AcquireReader() {
|
||||
t.Error("second AcquireReader must fail while the first holds it")
|
||||
}
|
||||
s.ReleaseReader()
|
||||
if !s.AcquireReader() {
|
||||
t.Error("AcquireReader should succeed again after release")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// ErrClosed is returned by Read/Write after Abort.
|
||||
var ErrClosed = errors.New("relay: ring aborted")
|
||||
|
||||
// Ring is a bounded byte buffer with single-producer / single-consumer
|
||||
// blocking semantics. Plan §M5: 单 session 64 MiB.
|
||||
//
|
||||
// Internally it's a flat slice with head/tail wrap-around, guarded by a
|
||||
// sync.Cond. Both Read and Write accept ctx so handlers can abort fast on
|
||||
// client disconnect — the watchdog goroutine broadcasts the cond when ctx
|
||||
// fires so blocked Wait calls re-check ctx.Err().
|
||||
type Ring struct {
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
|
||||
buf []byte
|
||||
head int
|
||||
tail int
|
||||
used int
|
||||
cap int
|
||||
|
||||
closed bool // sender called CloseWriter; reader drains to EOF
|
||||
aborted bool // both sides torn down
|
||||
}
|
||||
|
||||
func newRing(capacity int) *Ring {
|
||||
r := &Ring{
|
||||
buf: make([]byte, capacity),
|
||||
cap: capacity,
|
||||
}
|
||||
r.cond = sync.NewCond(&r.mu)
|
||||
return r
|
||||
}
|
||||
|
||||
func (r *Ring) Write(ctx context.Context, p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
stop := r.watchCtx(ctx)
|
||||
defer stop()
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
written := 0
|
||||
for written < len(p) {
|
||||
for r.used == r.cap && !r.aborted && !r.closed && ctxOK(ctx) {
|
||||
r.cond.Wait()
|
||||
}
|
||||
if !ctxOK(ctx) {
|
||||
return written, ctx.Err()
|
||||
}
|
||||
if r.aborted {
|
||||
return written, ErrClosed
|
||||
}
|
||||
if r.closed {
|
||||
return written, ErrClosed
|
||||
}
|
||||
|
||||
free := r.cap - r.used
|
||||
toWrite := len(p) - written
|
||||
if toWrite > free {
|
||||
toWrite = free
|
||||
}
|
||||
|
||||
end := r.tail + toWrite
|
||||
if end <= r.cap {
|
||||
copy(r.buf[r.tail:end], p[written:written+toWrite])
|
||||
} else {
|
||||
first := r.cap - r.tail
|
||||
copy(r.buf[r.tail:], p[written:written+first])
|
||||
copy(r.buf[:toWrite-first], p[written+first:written+toWrite])
|
||||
}
|
||||
r.tail = (r.tail + toWrite) % r.cap
|
||||
r.used += toWrite
|
||||
written += toWrite
|
||||
r.cond.Broadcast()
|
||||
}
|
||||
return written, nil
|
||||
}
|
||||
|
||||
func (r *Ring) Read(ctx context.Context, p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
stop := r.watchCtx(ctx)
|
||||
defer stop()
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
for r.used == 0 && !r.closed && !r.aborted && ctxOK(ctx) {
|
||||
r.cond.Wait()
|
||||
}
|
||||
if !ctxOK(ctx) {
|
||||
return 0, ctx.Err()
|
||||
}
|
||||
if r.used == 0 {
|
||||
if r.aborted {
|
||||
return 0, ErrClosed
|
||||
}
|
||||
return 0, io.EOF
|
||||
}
|
||||
|
||||
n := r.used
|
||||
if n > len(p) {
|
||||
n = len(p)
|
||||
}
|
||||
end := r.head + n
|
||||
if end <= r.cap {
|
||||
copy(p[:n], r.buf[r.head:end])
|
||||
} else {
|
||||
first := r.cap - r.head
|
||||
copy(p[:first], r.buf[r.head:])
|
||||
copy(p[first:n], r.buf[:n-first])
|
||||
}
|
||||
r.head = (r.head + n) % r.cap
|
||||
r.used -= n
|
||||
r.cond.Broadcast()
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// CloseWriter signals end-of-stream from the sender side. Reader returns
|
||||
// io.EOF after draining whatever is still buffered.
|
||||
func (r *Ring) CloseWriter() {
|
||||
r.mu.Lock()
|
||||
r.closed = true
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
// Abort tears the ring down hard. Subsequent Read/Write return ErrClosed.
|
||||
func (r *Ring) Abort() {
|
||||
r.mu.Lock()
|
||||
r.aborted = true
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *Ring) watchCtx(ctx context.Context) func() {
|
||||
if ctx == nil || ctx.Err() != nil {
|
||||
return func() {}
|
||||
}
|
||||
stop := make(chan struct{})
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
r.mu.Lock()
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
case <-stop:
|
||||
}
|
||||
}()
|
||||
return func() { close(stop) }
|
||||
}
|
||||
|
||||
func ctxOK(ctx context.Context) bool {
|
||||
if ctx == nil {
|
||||
return true
|
||||
}
|
||||
return ctx.Err() == nil
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRing_BasicWriteRead(t *testing.T) {
|
||||
r := newRing(64)
|
||||
want := []byte("hello, world")
|
||||
if n, err := r.Write(context.Background(), want); err != nil || n != len(want) {
|
||||
t.Fatalf("write: n=%d err=%v", n, err)
|
||||
}
|
||||
got := make([]byte, 32)
|
||||
n, err := r.Read(context.Background(), got)
|
||||
if err != nil {
|
||||
t.Fatalf("read: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got[:n], want) {
|
||||
t.Errorf("read mismatch: got %q want %q", got[:n], want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRing_CloseWriterSignalsEOF(t *testing.T) {
|
||||
r := newRing(32)
|
||||
_, _ = r.Write(context.Background(), []byte("abc"))
|
||||
r.CloseWriter()
|
||||
|
||||
got := make([]byte, 16)
|
||||
n, err := r.Read(context.Background(), got)
|
||||
if err != nil || n != 3 {
|
||||
t.Fatalf("first read: n=%d err=%v", n, err)
|
||||
}
|
||||
n, err = r.Read(context.Background(), got)
|
||||
if !errors.Is(err, io.EOF) {
|
||||
t.Errorf("expected EOF after drain; got n=%d err=%v", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRing_AbortReturnsErrClosed(t *testing.T) {
|
||||
r := newRing(32)
|
||||
r.Abort()
|
||||
if _, err := r.Write(context.Background(), []byte("x")); !errors.Is(err, ErrClosed) {
|
||||
t.Errorf("write after abort: got %v want ErrClosed", err)
|
||||
}
|
||||
if _, err := r.Read(context.Background(), make([]byte, 1)); !errors.Is(err, ErrClosed) {
|
||||
t.Errorf("read after abort: got %v want ErrClosed", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRing_BlocksWhenFullThenUnblocksOnRead(t *testing.T) {
|
||||
r := newRing(8)
|
||||
if _, err := r.Write(context.Background(), []byte("12345678")); err != nil {
|
||||
t.Fatalf("fill: %v", err)
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
_, err := r.Write(context.Background(), []byte("X"))
|
||||
if err != nil {
|
||||
t.Errorf("blocked write: %v", err)
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
|
||||
// Without a reader the goroutine should block.
|
||||
select {
|
||||
case <-done:
|
||||
t.Fatal("write should block when ring is full")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
|
||||
got := make([]byte, 4)
|
||||
if _, err := r.Read(context.Background(), got); err != nil {
|
||||
t.Fatalf("read: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("write did not unblock after read")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRing_CtxCancelsBlockedRead(t *testing.T) {
|
||||
r := newRing(8)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
type res struct {
|
||||
n int
|
||||
err error
|
||||
}
|
||||
out := make(chan res, 1)
|
||||
go func() {
|
||||
buf := make([]byte, 4)
|
||||
n, err := r.Read(ctx, buf)
|
||||
out <- res{n, err}
|
||||
}()
|
||||
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case r := <-out:
|
||||
if !errors.Is(r.err, context.Canceled) {
|
||||
t.Errorf("expected context.Canceled, got %v", r.err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("blocked read did not unblock on ctx cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRing_WrapsAround(t *testing.T) {
|
||||
r := newRing(4)
|
||||
// Cycle through enough bytes to force head/tail wrapping.
|
||||
for i := 0; i < 20; i++ {
|
||||
_, _ = r.Write(context.Background(), []byte{byte(i)})
|
||||
got := make([]byte, 1)
|
||||
n, _ := r.Read(context.Background(), got)
|
||||
if n != 1 || got[0] != byte(i) {
|
||||
t.Errorf("iter %d: got %v want [%d]", i, got, i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRing_StressManyChunks(t *testing.T) {
|
||||
const total = 4 * 1024
|
||||
r := newRing(64)
|
||||
|
||||
var wgWriter, wgReader sync.WaitGroup
|
||||
wgWriter.Add(1)
|
||||
wgReader.Add(1)
|
||||
|
||||
src := make([]byte, total)
|
||||
for i := range src {
|
||||
src[i] = byte(i)
|
||||
}
|
||||
|
||||
go func() {
|
||||
defer wgWriter.Done()
|
||||
for i := 0; i < total; i += 13 {
|
||||
end := i + 13
|
||||
if end > total {
|
||||
end = total
|
||||
}
|
||||
_, _ = r.Write(context.Background(), src[i:end])
|
||||
}
|
||||
r.CloseWriter()
|
||||
}()
|
||||
|
||||
got := make([]byte, 0, total)
|
||||
go func() {
|
||||
defer wgReader.Done()
|
||||
buf := make([]byte, 17)
|
||||
for {
|
||||
n, err := r.Read(context.Background(), buf)
|
||||
got = append(got, buf[:n]...)
|
||||
if errors.Is(err, io.EOF) {
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Errorf("read: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
wgWriter.Wait()
|
||||
wgReader.Wait()
|
||||
if !bytes.Equal(got, src) {
|
||||
t.Errorf("mismatch: got %d bytes, want %d", len(got), len(src))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user