rpc/conn.go

   1 // Simple RPC-like protocol: implementation of connection and basic actions
   2
   3 // SPDX-License-Identifier: GPL-3.0-or-later
   4 // Copyright (C) 2021-2024  Simon Ruderich
   5
   6 package rpc
   7
   8 import (
   9         "bufio"
  10         "fmt"
  11         "io/fs"
  12         "os/exec"
  13         "strings"
  14         "sync"
  15
  16         "ruderich.org/simon/safcm"
  17 )
  18
  19 type Conn struct {
  20         Events   <-chan ConnEvent
  21         events   chan<- ConnEvent // same as Events, to publish events
  22         eventsWg sync.WaitGroup
  23
  24         debug     bool
  25         sshRemote string
  26         sshOpts   []string
  27
  28         remoteHelpers fs.FS
  29
  30         cmd  *exec.Cmd
  31         conn *safcm.GobConn
  32 }
  33
  34 type ConnEventType int
  35
  36 const (
  37         _               ConnEventType = iota
  38         ConnEventStderr               // stderr from spawned process
  39         ConnEventDebug                // debug message
  40         ConnEventUpload               // remote helper upload in progress
  41 )
  42
  43 type ConnEvent struct {
  44         Type ConnEventType
  45         Data string
  46 }
  47
  48 // NewConn creates a new connection. Events in the returned struct must be
  49 // regularly read or the connection will hang. This must be done before
  50 // DialSSH is called to open a connection.
  51 func NewConn(debug bool) *Conn {
  52         ch := make(chan ConnEvent)
  53         return &Conn{
  54                 Events: ch,
  55                 events: ch,
  56                 debug:  debug,
  57         }
  58 }
  59
  60 func (c *Conn) debugf(format string, a ...interface{}) {
  61         if !c.debug {
  62                 return
  63         }
  64         c.events <- ConnEvent{
  65                 Type: ConnEventDebug,
  66                 Data: fmt.Sprintf(format, a...),
  67         }
  68 }
  69
  70 // Wrap safcm.GobConn's Send() and Recv() to provide debug output.
  71
  72 // Send sends a single message to the remote.
  73 func (c *Conn) Send(m safcm.Msg) error {
  74         // No checks for invalid Conn, a stacktrace is more helpful
  75
  76         c.debugf("Send: sending %#v", m)
  77         return c.conn.Send(m)
  78 }
  79
  80 // Recv waits for a single message from the remote.
  81 func (c *Conn) Recv() (safcm.Msg, error) {
  82         // No checks for invalid Conn, a stacktrace is more helpful
  83
  84         c.debugf("Recv: waiting for message")
  85         m, err := c.conn.Recv()
  86         c.debugf("Recv: received msg=%#v err=%#v", m, err)
  87         return m, err
  88 }
  89
  90 // Wait waits for the connection to terminate. It's safe to call Wait (and
  91 // Kill) multiple times.
  92 func (c *Conn) Wait() error {
  93         // But check here because Wait() can be called multiple times
  94         if c.cmd == nil {
  95                 return fmt.Errorf("Dial*() not called or already terminated")
  96         }
  97
  98         c.debugf("Wait: waiting for connection to terminate")
  99         return c.wait()
 100 }
 101 func (c *Conn) wait() error {
 102         err := c.cmd.Wait()
 103         c.cmd = nil
 104
 105         // Wait until we've received all events from the program's stderr.
 106         c.eventsWg.Wait()
 107         // Notify consumers that no more events will occur.
 108         close(c.events)
 109         // We cannot reuse this channel.
 110         c.events = nil
 111         // Don't set c.Events to nil because this creates a data race when
 112         // another thread is still waiting.
 113
 114         return err
 115 }
 116
 117 // Kill forcefully terminates the connection. It's safe to call Kill (and
 118 // Wait) multiple times. Calling it before Dial*() was called will only close
 119 // the Events channel.
 120 func (c *Conn) Kill() error {
 121         if c.cmd == nil {
 122                 if c.events != nil {
 123                         close(c.events)
 124                         c.events = nil
 125                 }
 126                 return fmt.Errorf("Dial*() not called or already terminated")
 127         }
 128
 129         c.debugf("Kill: killing connection")
 130
 131         c.cmd.Process.Kill() //nolint:errcheck
 132         return c.wait()
 133 }
 134
 135 func (c *Conn) handleStderrAsEvents(cmd *exec.Cmd) error {
 136         // cmd may differ from c.cmd here!
 137         stderr, err := cmd.StderrPipe()
 138         if err != nil {
 139                 return err
 140         }
 141
 142         c.eventsWg.Add(1)
 143         go func() {
 144                 r := bufio.NewReader(stderr)
 145                 for {
 146                         x, err := r.ReadString('\n')
 147                         if err != nil {
 148                                 break
 149                         }
 150                         x = strings.TrimRight(x, "\n")
 151
 152                         c.events <- ConnEvent{
 153                                 Type: ConnEventStderr,
 154                                 Data: x,
 155                         }
 156                 }
 157                 c.eventsWg.Done()
 158         }()
 159
 160         return nil
 161 }