diff --git a/decrypt.go b/decrypt.go index 84d627d..630e963 100644 --- a/decrypt.go +++ b/decrypt.go @@ -66,6 +66,7 @@ func (d *Dispatcher) AddInstance(inst *WrapperInstance) { return } decryptInstance.onCapacity = d.signalCapacity + decryptInstance.onUnavailable = d.quarantineInstance var replaced *DecryptInstance d.mu.Lock() @@ -90,6 +91,26 @@ func (d *Dispatcher) AddInstance(inst *WrapperInstance) { logrus.Debugf("added instance %s", inst.Id) } +func (d *Dispatcher) quarantineInstance(target *DecryptInstance, _ string) { + if target == nil { + return + } + removed := false + d.mu.Lock() + for i, current := range d.Instances { + if current == target { + d.generation[target.id]++ + d.Instances = append(d.Instances[:i], d.Instances[i+1:]...) + removed = true + break + } + } + d.mu.Unlock() + if removed { + d.signalCapacity() + } +} + func (d *Dispatcher) RemoveInstance(id string) { var removed *DecryptInstance d.mu.Lock() diff --git a/decrypt_instance.go b/decrypt_instance.go index 10ba28a..737fbf6 100644 --- a/decrypt_instance.go +++ b/decrypt_instance.go @@ -14,10 +14,14 @@ import ( ) const ( - defaultId = "0" - prefetchKey = "skd://itunes.apple.com/P000000000/s1/e1" - decryptIOTimeout = 30 * time.Second - maxPoolSize = 10 + defaultId = "0" + prefetchKey = "skd://itunes.apple.com/P000000000/s1/e1" + decryptIOTimeout = 30 * time.Second + maxPoolSize = 10 + wrapperFailureWindow = 60 * time.Second + wrapperFailureThreshold = 3 + wrapperFailureMinConns = 3 + wrapperFailureMinAdamIDs = 2 ) var errInstanceBusy = errors.New("decrypt instance is at capacity") @@ -33,6 +37,12 @@ type decryptConn struct { writeBuffers net.Buffers } +type wrapperIOFailure struct { + at time.Time + conn *decryptConn + adamID string +} + // DecryptSession leases one wrapper connection for the lifetime of a client // gRPC stream. The stream context cancels pool waits, dials, and blocked I/O. type DecryptSession struct { @@ -56,30 +66,39 @@ type DecryptInstance struct { region string decryptPort int - poolMu sync.Mutex - pool []*decryptConn - connections map[*decryptConn]struct{} - reserved int - isClosed bool - poolLimit int - dialContext dialContextFunc - ioTimeout time.Duration - onCapacity func() - - closeOnce sync.Once + poolMu sync.Mutex + pool []*decryptConn + connections map[*decryptConn]struct{} + reserved int + isClosed bool + poolLimit int + dialContext dialContextFunc + ioTimeout time.Duration + onCapacity func() + onUnavailable func(*DecryptInstance, string) + terminateWrapper func() error + now func() time.Time + + healthMu sync.Mutex + failures []wrapperIOFailure + + closeOnce sync.Once + unavailableOnce sync.Once } func NewDecryptInstance(inst *WrapperInstance) (*DecryptInstance, error) { dialer := &net.Dialer{Timeout: 10 * time.Second} instance := &DecryptInstance{ - id: inst.Id, - region: inst.Region, - decryptPort: inst.DecryptPort, - pool: make([]*decryptConn, 0, maxPoolSize), - connections: make(map[*decryptConn]struct{}, maxPoolSize), - poolLimit: maxPoolSize, - dialContext: dialer.DialContext, - ioTimeout: decryptIOTimeout, + id: inst.Id, + region: inst.Region, + decryptPort: inst.DecryptPort, + pool: make([]*decryptConn, 0, maxPoolSize), + connections: make(map[*decryptConn]struct{}, maxPoolSize), + poolLimit: maxPoolSize, + dialContext: dialer.DialContext, + ioTimeout: decryptIOTimeout, + terminateWrapper: func() error { return terminateWrapperInstance(inst, wrapperTerminateGrace) }, + now: time.Now, } // Pre-warm one connection both to validate the wrapper and to keep the @@ -271,11 +290,28 @@ func (d *DecryptInstance) Close() { }) } -func (d *DecryptInstance) Unavailable() { - d.Close() - if err := KillWrapper(d.id); err != nil { - logrus.Errorf("failed to kill instance %s: %s", d.id, err) - } +func (d *DecryptInstance) Unavailable(reason string) { + d.unavailableOnce.Do(func() { + // Closing first immediately removes this instance from scheduling and + // interrupts every leased connection. The wrapper lifecycle will replace + // the process and register a fresh DecryptInstance after it exits. + d.Close() + logrus.Warnf("wrapper instance %s is unhealthy: %s; restarting", d.id, reason) + if d.onUnavailable != nil { + d.onUnavailable(d, reason) + } + if d.terminateWrapper == nil { + logrus.Errorf("failed to restart instance %s: no wrapper kill function", d.id) + return + } + // Process termination may wait for a grace period. It must not delay the + // failed decrypt response or hold up healthy instances in the dispatcher. + go func() { + if err := d.terminateWrapper(); err != nil { + logrus.Errorf("failed to terminate instance %s: %s", d.id, err) + } + }() + }) } func (s *DecryptSession) currentConn() (*decryptConn, error) { @@ -298,6 +334,7 @@ func (s *DecryptSession) Decrypt(adamId, key string, payload []byte) ([]byte, er } if c.lastAdamId != adamId || c.lastKey != key { if err := s.instance.switchConnContext(s.ctx, c, adamId, key); err != nil { + s.instance.observeWrapperIOFailure(s.ctx, c, adamId, "context switch", err) s.Discard() return nil, mapContextError(s.ctx, err) } @@ -307,12 +344,96 @@ func (s *DecryptSession) Decrypt(adamId, key string, payload []byte) ([]byte, er // The slice is sent once and is never modified by a later request. result, err := s.instance.decryptConn(s.ctx, c, payload, payload) if err != nil { + s.instance.observeWrapperIOFailure(s.ctx, c, adamId, "decrypt", err) s.Discard() return nil, mapContextError(s.ctx, err) } return result, nil } +func classifyLocalWrapperIOError(ctx context.Context, err error) (local, timedOut bool) { + if err == nil { + return false, false + } + // A client cancellation or client-owned deadline is not evidence that the + // local wrapper process is unhealthy. + if ctx != nil && ctx.Err() != nil { + return false, false + } + if ctx != nil { + if deadline, ok := ctx.Deadline(); ok && !time.Now().Before(deadline) { + return false, false + } + } + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, net.ErrClosed) { + return true, false + } + var netErr net.Error + if !errors.As(err, &netErr) { + return false, false + } + return true, netErr.Timeout() +} + +func (d *DecryptInstance) observeWrapperIOFailure(ctx context.Context, conn *decryptConn, adamID, stage string, err error) { + local, timedOut := classifyLocalWrapperIOError(ctx, err) + if conn == nil || adamID == "" || !local { + return + } + d.poolMu.Lock() + closed := d.isClosed + d.poolMu.Unlock() + if closed { + return + } + // The deadline owned by this manager is thirty seconds. A loopback wrapper + // operation exceeding it is already conclusive evidence of a wedged local + // process; waiting for more timeouts only prolongs the outage, especially at + // concurrency one. + if timedOut { + d.Unavailable(fmt.Sprintf("local %s I/O timed out for Adam ID %s", stage, adamID)) + return + } + + now := time.Now() + if d.now != nil { + now = d.now() + } + cutoff := now.Add(-wrapperFailureWindow) + + d.healthMu.Lock() + kept := d.failures[:0] + for _, failure := range d.failures { + if !failure.at.Before(cutoff) { + kept = append(kept, failure) + } + } + d.failures = append(kept, wrapperIOFailure{ + at: now, + conn: conn, + adamID: adamID, + }) + + connections := make(map[*decryptConn]struct{}, len(d.failures)) + adamIDs := make(map[string]struct{}, len(d.failures)) + for _, failure := range d.failures { + connections[failure.conn] = struct{}{} + adamIDs[failure.adamID] = struct{}{} + } + failureCount := len(d.failures) + shouldTrip := failureCount >= wrapperFailureThreshold && len(connections) >= wrapperFailureMinConns && len(adamIDs) >= wrapperFailureMinAdamIDs + d.healthMu.Unlock() + + if !shouldTrip { + logrus.Warnf("wrapper instance %s local %s I/O failure (%d/%d in %s): %v", d.id, stage, failureCount, wrapperFailureThreshold, wrapperFailureWindow, err) + return + } + d.Unavailable(fmt.Sprintf( + "%d local I/O failures across %d connections and %d Adam IDs in %s", + failureCount, len(connections), len(adamIDs), wrapperFailureWindow, + )) +} + func mapContextError(ctx context.Context, err error) error { if ctxErr := ctx.Err(); ctxErr != nil { return ctxErr diff --git a/decrypt_instance_test.go b/decrypt_instance_test.go index b764a59..fd763f6 100644 --- a/decrypt_instance_test.go +++ b/decrypt_instance_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/binary" "errors" + "fmt" "io" "net" "sync" @@ -22,6 +23,12 @@ type fakeDecryptServer struct { conns map[net.Conn]struct{} } +type fakeTimeoutError struct{} + +func (fakeTimeoutError) Error() string { return "i/o timeout" } +func (fakeTimeoutError) Timeout() bool { return true } +func (fakeTimeoutError) Temporary() bool { return true } + func newFakeDecryptServer(t *testing.T) *fakeDecryptServer { t.Helper() listener, err := net.Listen("tcp", "127.0.0.1:0") @@ -373,6 +380,166 @@ func TestDecryptCancellationWithoutDeadlineInterruptsBlockedWrapperIO(t *testing } } +func TestWrapperIOTimeoutQuarantinesAndTerminatesInstanceOnce(t *testing.T) { + server := newFakeDecryptServer(t) + instance, err := NewDecryptInstance(&WrapperInstance{Id: "test", Region: "cn", DecryptPort: server.port()}) + if err != nil { + t.Fatal(err) + } + instance.ioTimeout = 30 * time.Millisecond + terminated := make(chan struct{}, 2) + instance.terminateWrapper = func() error { + terminated <- struct{}{} + return nil + } + + session, err := instance.OpenSession(context.Background(), "song-1", "key-1") + if err != nil { + t.Fatal(err) + } + if _, err := session.Decrypt("song-1", "key-1", []byte("block")); err == nil { + t.Fatal("expected wrapper I/O timeout") + } + select { + case <-terminated: + case <-time.After(time.Second): + t.Fatal("timed out waiting for wrapper termination") + } + instance.Unavailable("duplicate trigger") + select { + case <-terminated: + t.Fatal("wrapper termination ran more than once") + case <-time.After(50 * time.Millisecond): + } + instance.poolMu.Lock() + defer instance.poolMu.Unlock() + if !instance.isClosed || len(instance.connections) != 0 || len(instance.pool) != 0 { + t.Fatalf("unhealthy instance was not quarantined: closed=%v connections=%d pool=%d", instance.isClosed, len(instance.connections), len(instance.pool)) + } +} + +func TestConcurrentWrapperIOTimeoutsScheduleSingleTermination(t *testing.T) { + instance := &DecryptInstance{ + id: "test", + connections: make(map[*decryptConn]struct{}), + poolLimit: maxPoolSize, + } + var terminations atomic.Int32 + terminationStarted := make(chan struct{}) + releaseTermination := make(chan struct{}) + instance.terminateWrapper = func() error { + if terminations.Add(1) == 1 { + close(terminationStarted) + } + <-releaseTermination + return nil + } + + start := make(chan struct{}) + var wg sync.WaitGroup + for i := 0; i < maxPoolSize; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + instance.observeWrapperIOFailure( + context.Background(), + &decryptConn{}, + fmt.Sprintf("song-%d", i), + "decrypt", + fakeTimeoutError{}, + ) + }(i) + } + close(start) + returned := make(chan struct{}) + go func() { + wg.Wait() + close(returned) + }() + select { + case <-terminationStarted: + case <-time.After(time.Second): + t.Fatal("wrapper termination did not start") + } + select { + case <-returned: + case <-time.After(time.Second): + t.Fatal("failure reporters blocked on wrapper termination") + } + if got := terminations.Load(); got != 1 { + t.Fatalf("wrapper terminations = %d, want 1", got) + } + close(releaseTermination) +} + +func TestClientDeadlineDoesNotQuarantineWrapper(t *testing.T) { + server := newFakeDecryptServer(t) + instance, err := NewDecryptInstance(&WrapperInstance{Id: "test", Region: "cn", DecryptPort: server.port()}) + if err != nil { + t.Fatal(err) + } + instance.ioTimeout = time.Second + var terminations atomic.Int32 + instance.terminateWrapper = func() error { + terminations.Add(1) + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond) + defer cancel() + session, err := instance.OpenSession(ctx, "song-1", "key-1") + if err != nil { + t.Fatal(err) + } + if _, err := session.Decrypt("song-1", "key-1", []byte("block")); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("Decrypt error = %v, want deadline exceeded", err) + } + time.Sleep(20 * time.Millisecond) + if got := terminations.Load(); got != 0 { + t.Fatalf("client deadline triggered %d wrapper terminations", got) + } + instance.poolMu.Lock() + defer instance.poolMu.Unlock() + if instance.isClosed { + t.Fatal("client deadline quarantined wrapper instance") + } +} + +func TestRepeatedConnectionFailuresAcrossSongsQuarantineOnce(t *testing.T) { + instance := &DecryptInstance{ + id: "test", + connections: make(map[*decryptConn]struct{}), + poolLimit: maxPoolSize, + now: time.Now, + } + terminated := make(chan struct{}, 2) + instance.terminateWrapper = func() error { + terminated <- struct{}{} + return nil + } + for i, adamID := range []string{"song-1", "song-1", "song-2"} { + instance.observeWrapperIOFailure(context.Background(), &decryptConn{}, adamID, "decrypt", io.EOF) + if i < 2 { + select { + case <-terminated: + t.Fatalf("wrapper terminated after only %d connection failures", i+1) + default: + } + } + } + select { + case <-terminated: + case <-time.After(time.Second): + t.Fatal("repeated cross-song connection failures did not terminate wrapper") + } + instance.Unavailable("duplicate trigger") + select { + case <-terminated: + t.Fatal("wrapper termination ran more than once") + case <-time.After(50 * time.Millisecond): + } +} + func TestDecryptPartialFailureDiscardsConnection(t *testing.T) { server := newFakeDecryptServer(t) instance, err := NewDecryptInstance(&WrapperInstance{Id: "test", Region: "cn", DecryptPort: server.port()}) diff --git a/decrypt_test.go b/decrypt_test.go index 867542d..247209b 100644 --- a/decrypt_test.go +++ b/decrypt_test.go @@ -146,6 +146,49 @@ func TestDispatcherRoutesAroundFullInstance(t *testing.T) { } } +func TestDispatcherRoutesAroundQuarantinedInstance(t *testing.T) { + d, instances := newTestDispatcher(t, 2, 2) + bad := instances[0] + bad.onUnavailable = d.quarantineInstance + terminated := make(chan struct{}, 1) + bad.terminateWrapper = func() error { + terminated <- struct{}{} + return nil + } + bad.Unavailable("test quarantine") + select { + case <-terminated: + case <-time.After(time.Second): + t.Fatal("timed out waiting for wrapper termination") + } + + for i := 0; i < 2; i++ { + session, err := d.OpenSession(context.Background(), "song", "key") + if err != nil { + t.Fatal(err) + } + if session.instance != instances[1] { + t.Fatal("dispatcher selected quarantined instance") + } + session.Close() + } + if got := d.snapshotInstances(); len(got) != 1 || got[0] != instances[1] { + t.Fatalf("dispatcher instances = %#v, want only healthy instance", got) + } +} + +func TestQuarantineDoesNotRemoveSameIDReplacement(t *testing.T) { + d := NewDispatcher() + old := &DecryptInstance{id: "same", connections: make(map[*decryptConn]struct{})} + replacement := &DecryptInstance{id: "same", connections: make(map[*decryptConn]struct{})} + d.Instances = []*DecryptInstance{replacement} + d.quarantineInstance(old, "stale failure") + got := d.snapshotInstances() + if len(got) != 1 || got[0] != replacement { + t.Fatal("stale quarantine removed same-ID replacement") + } +} + func TestDispatcherFullPoolsHonorCancellation(t *testing.T) { d, instances := newTestDispatcher(t, 2, 1) for _, instance := range instances { diff --git a/instance.go b/instance.go index b93335f..d8024c8 100644 --- a/instance.go +++ b/instance.go @@ -4,18 +4,22 @@ import ( "encoding/json" "os" "os/exec" + "time" ) +const wrapperTerminateGrace = 5 * time.Second + var Instances []*WrapperInstance type WrapperInstance struct { - Id string `json:"id"` - Account string `json:"account"` - Region string `json:"region"` - DecryptPort int `json:"-"` - M3U8Port int `json:"-"` - NoRestart bool `json:"-"` - Cmd *exec.Cmd `json:"-"` + Id string `json:"id"` + Account string `json:"account"` + Region string `json:"region"` + DecryptPort int `json:"-"` + M3U8Port int `json:"-"` + NoRestart bool `json:"-"` + Cmd *exec.Cmd `json:"-"` + Done chan struct{} `json:"-"` } func SaveInstances() { diff --git a/port.go b/port.go index 4e3e0dc..77ac110 100644 --- a/port.go +++ b/port.go @@ -47,3 +47,12 @@ func GenerateUniquePort() int { return port } } + +func ReleasePort(port int) { + if port < 0 { + return + } + portMutex.Lock() + delete(usedPorts, port) + portMutex.Unlock() +} diff --git a/wrapper.go b/wrapper.go index 7337255..4640fe6 100644 --- a/wrapper.go +++ b/wrapper.go @@ -3,6 +3,7 @@ package main import ( "bufio" "encoding/json" + "errors" "fmt" "github.com/artdarek/go-unzip" "github.com/creack/pty" @@ -14,6 +15,7 @@ import ( "runtime" "strconv" "strings" + "time" ) type wrapperRelease struct { @@ -114,6 +116,7 @@ func WrapperInitial(id uuid.UUID, account string, password string) { DecryptPort: GenerateUniquePort(), M3U8Port: GenerateUniquePort(), NoRestart: true, + Done: make(chan struct{}), } args := []string{ @@ -143,6 +146,7 @@ func WrapperInitial(id uuid.UUID, account string, password string) { go handleOutput(ptmx, &instance) err = cmd.Wait() + close(instance.Done) if err != nil { log.Warnf("Wrapper exited with error: %v\n", err) } @@ -163,6 +167,7 @@ func WrapperStart(id string, account string) { DecryptPort: GenerateUniquePort(), M3U8Port: GenerateUniquePort(), NoRestart: false, + Done: make(chan struct{}), } args := []string{ @@ -190,6 +195,7 @@ func WrapperStart(id string, account string) { go handleOutput(ptmx, &instance) _ = cmd.Wait() + close(instance.Done) go wrapperDown(&instance) } @@ -236,6 +242,8 @@ func wrapperReady(instance *WrapperInstance) { func wrapperDown(instance *WrapperInstance) { log.Info(fmt.Sprintf("[wrapper %s]", strings.Split(instance.Id, "-")[0]), " Wrapper Down") + ReleasePort(instance.DecryptPort) + ReleasePort(instance.M3U8Port) RemoveInstance(instance) WMDispatcher.RemoveInstance(instance.Id) if !instance.NoRestart { @@ -247,14 +255,21 @@ func wrapperDown(instance *WrapperInstance) { func KillWrapper(id string) error { instance := GetInstance(id) - if instance == nil { + if instance == nil || instance.Id == "" { return fmt.Errorf("instance %s not found", id) } + return signalWrapperInstance(instance) +} + +func signalWrapperInstance(instance *WrapperInstance) error { + if instance == nil { + return fmt.Errorf("wrapper instance is nil") + } if instance.Cmd == nil { - return fmt.Errorf("instance %s cmd is nil", id) + return fmt.Errorf("instance %s cmd is nil", instance.Id) } if instance.Cmd.Process == nil { - return fmt.Errorf("instance %s process is nil", id) + return fmt.Errorf("instance %s process is nil", instance.Id) } // Send SIGINT to trigger wrapper's internal child-killing signal handler err := instance.Cmd.Process.Signal(os.Interrupt) @@ -264,6 +279,36 @@ func KillWrapper(id string) error { return nil } +func terminateWrapperInstance(instance *WrapperInstance, grace time.Duration) error { + if err := signalWrapperInstance(instance); err != nil { + if errors.Is(err, os.ErrProcessDone) { + return nil + } + return err + } + if instance.Done == nil { + return nil + } + timer := time.NewTimer(grace) + defer timer.Stop() + select { + case <-instance.Done: + return nil + case <-timer.C: + } + log.Warnf("[wrapper %s] did not exit within %s after SIGINT; sending SIGKILL", strings.Split(instance.Id, "-")[0], grace) + if err := instance.Cmd.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) { + return err + } + timer.Reset(grace) + select { + case <-instance.Done: + return nil + case <-timer.C: + return fmt.Errorf("instance %s did not exit after SIGKILL", instance.Id) + } +} + func provide2FACode(id string, code string) { path := "data/wrapper/rootfs/data/instances/" + id + "/2fa.txt" err := os.WriteFile(path, []byte(code), 0777) diff --git a/wrapper_lifecycle_test.go b/wrapper_lifecycle_test.go new file mode 100644 index 0000000..348cf08 --- /dev/null +++ b/wrapper_lifecycle_test.go @@ -0,0 +1,87 @@ +package main + +import ( + "bufio" + "fmt" + "os" + "os/exec" + "os/signal" + "testing" + "time" +) + +func TestWrapperProcessHelper(t *testing.T) { + mode := os.Getenv("WRAPPER_MANAGER_HELPER_MODE") + if mode == "" { + return + } + if mode == "ignore" { + signal.Ignore(os.Interrupt) + fmt.Println("ready") + select {} + } + interrupt := make(chan os.Signal, 1) + signal.Notify(interrupt, os.Interrupt) + fmt.Println("ready") + <-interrupt +} + +func startWrapperProcessHelper(t *testing.T, mode string) *WrapperInstance { + t.Helper() + cmd := exec.Command(os.Args[0], "-test.run=TestWrapperProcessHelper") + cmd.Env = append(os.Environ(), "WRAPPER_MANAGER_HELPER_MODE="+mode) + stdout, err := cmd.StdoutPipe() + if err != nil { + t.Fatal(err) + } + if err := cmd.Start(); err != nil { + t.Fatal(err) + } + if scanner := bufio.NewScanner(stdout); !scanner.Scan() || scanner.Text() != "ready" { + _ = cmd.Process.Kill() + t.Fatal("wrapper process helper did not become ready") + } + done := make(chan struct{}) + go func() { + _ = cmd.Wait() + close(done) + }() + t.Cleanup(func() { + if cmd.ProcessState == nil || !cmd.ProcessState.Exited() { + _ = cmd.Process.Kill() + } + select { + case <-done: + case <-time.After(time.Second): + } + }) + return &WrapperInstance{Id: "helper", Cmd: cmd, Done: done} +} + +func TestTerminateHungWrapperEscalatesToKill(t *testing.T) { + instance := startWrapperProcessHelper(t, "ignore") + started := time.Now() + if err := terminateWrapperInstance(instance, 50*time.Millisecond); err != nil { + t.Fatal(err) + } + if elapsed := time.Since(started); elapsed < 50*time.Millisecond { + t.Fatalf("hung wrapper terminated before grace period elapsed: %s", elapsed) + } + select { + case <-instance.Done: + case <-time.After(time.Second): + t.Fatal("hung wrapper process did not exit after SIGKILL") + } +} + +func TestTerminateResponsiveWrapperExitsDuringGrace(t *testing.T) { + instance := startWrapperProcessHelper(t, "responsive") + if err := terminateWrapperInstance(instance, time.Second); err != nil { + t.Fatal(err) + } + select { + case <-instance.Done: + case <-time.After(time.Second): + t.Fatal("responsive wrapper process did not exit after SIGINT") + } +}