| 1 | package forward |
| 2 | |
| 3 | import ( |
| 4 | "errors" |
| 5 | "io" |
| 6 | "net" |
| 7 | "sync" |
| 8 | |
| 9 | "golang.org/x/crypto/ssh" |
| 10 | ) |
| 11 | |
| 12 | // Event reports a forward's transition, delivered to the Set's onEvent hook. |
| 13 | type Event struct { |
| 14 | Spec Spec |
| 15 | Up bool |
| 16 | BoundAddr string // actual bound address (resolves ":0") |
| 17 | Err error |
| 18 | } |
| 19 | |
| 20 | // Entry is a snapshot of one registered forward. |
| 21 | type Entry struct { |
| 22 | Spec Spec |
| 23 | Up bool |
| 24 | BoundAddr string |
| 25 | LastErr error |
| 26 | } |
| 27 | |
| 28 | // Set is the port-forward registry for one Client. It is safe for concurrent |
| 29 | // use. Local listeners persist across Detach/Attach so a forwarded port stays |
| 30 | // reserved through reconnects; remote listeners are torn down on Detach and |
| 31 | // recreated on Attach. |
| 32 | type Set struct { |
| 33 | onEvent func(Event) |
| 34 | |
| 35 | mu sync.Mutex |
| 36 | ssh *ssh.Client |
| 37 | runs map[string]*runner |
| 38 | } |
| 39 | |
| 40 | // NewSet creates an empty Set. onEvent may be nil. |
| 41 | func NewSet(onEvent func(Event)) *Set { |
| 42 | return &Set{onEvent: onEvent, runs: map[string]*runner{}} |
| 43 | } |
| 44 | |
| 45 | type runner struct { |
| 46 | spec Spec |
| 47 | local net.Listener // Local: persistent local listener |
| 48 | remote net.Listener // Remote: per-connection remote listener |
| 49 | up bool |
| 50 | lastErr error |
| 51 | stop chan struct{} |
| 52 | acceptWG sync.WaitGroup |
| 53 | } |
| 54 | |
| 55 | // Add registers spec and starts it if the Set is attached. Returns the bound |
| 56 | // address (useful for ":0" local forwards). |
| 57 | func (s *Set) Add(spec Spec) (string, error) { |
| 58 | if err := spec.Validate(); err != nil { |
| 59 | return "", err |
| 60 | } |
| 61 | name := spec.DefaultName() |
| 62 | spec.Name = name |
| 63 | s.mu.Lock() |
| 64 | defer s.mu.Unlock() |
| 65 | if _, ok := s.runs[name]; ok { |
| 66 | return "", ErrDuplicateForward |
| 67 | } |
| 68 | r := &runner{spec: spec, stop: make(chan struct{})} |
| 69 | s.runs[name] = r |
| 70 | if s.ssh == nil { |
| 71 | return "", nil // starts on Attach |
| 72 | } |
| 73 | bound, err := s.startLocked(r, s.ssh) |
| 74 | if err != nil { |
| 75 | delete(s.runs, name) |
| 76 | return "", err |
| 77 | } |
| 78 | return bound, nil |
| 79 | } |
| 80 | |
| 81 | // Replace atomically swaps the named registration after the replacement has |
| 82 | // started successfully. If startup fails, the existing forward remains live. |
| 83 | // This is primarily used when a remote serve moves to a new workspace/port. |
| 84 | func (s *Set) Replace(spec Spec) (string, error) { |
| 85 | if err := spec.Validate(); err != nil { |
| 86 | return "", err |
| 87 | } |
| 88 | name := spec.DefaultName() |
| 89 | spec.Name = name |
| 90 | s.mu.Lock() |
| 91 | old := s.runs[name] |
| 92 | if s.ssh == nil { |
| 93 | s.mu.Unlock() |
| 94 | return "", ErrNotAttached |
| 95 | } |
| 96 | replacement := &runner{spec: spec, stop: make(chan struct{})} |
| 97 | bound := "" |
| 98 | bound, err := s.startLocked(replacement, s.ssh) |
| 99 | if err != nil { |
| 100 | s.mu.Unlock() |
| 101 | return "", err |
| 102 | } |
| 103 | s.runs[name] = replacement |
| 104 | if old != nil { |
| 105 | // The replacement's Up event is authoritative; suppress a later Down event |
| 106 | // from retiring the old runner with the same name. |
| 107 | old.up = false |
| 108 | } |
| 109 | s.mu.Unlock() |
| 110 | if old != nil { |
| 111 | s.stopRunner(old, true) |
| 112 | } |
| 113 | return bound, nil |
| 114 | } |
| 115 | |
| 116 | // Remove stops and deregisters the named forward. |
| 117 | func (s *Set) Remove(name string) error { |
| 118 | s.mu.Lock() |
| 119 | r, ok := s.runs[name] |
| 120 | if ok { |
| 121 | delete(s.runs, name) |
| 122 | } |
| 123 | s.mu.Unlock() |
| 124 | if !ok { |
| 125 | return errors.New("forward: no such forward: " + name) |
| 126 | } |
| 127 | s.stopRunner(r, true) |
| 128 | return nil |
| 129 | } |
| 130 | |
| 131 | // List snapshots all registered forwards. |
| 132 | func (s *Set) List() []Entry { |
| 133 | s.mu.Lock() |
| 134 | defer s.mu.Unlock() |
| 135 | out := make([]Entry, 0, len(s.runs)) |
| 136 | for _, r := range s.runs { |
| 137 | bound := "" |
| 138 | if r.local != nil { |
| 139 | bound = r.local.Addr().String() |
| 140 | } else if r.remote != nil { |
| 141 | bound = r.remote.Addr().String() |
| 142 | } |
| 143 | out = append(out, Entry{Spec: r.spec, Up: r.up, BoundAddr: bound, LastErr: r.lastErr}) |
| 144 | } |
| 145 | return out |
| 146 | } |
| 147 | |
| 148 | // Attach binds the Set to a (re)connected ssh client and (re)starts every |
| 149 | // forward. Per-forward failures are joined and returned; successfully started |
| 150 | // forwards stay up. |
| 151 | func (s *Set) Attach(cl *ssh.Client) error { |
| 152 | s.mu.Lock() |
| 153 | defer s.mu.Unlock() |
| 154 | s.ssh = cl |
| 155 | var errs []error |
| 156 | for _, r := range s.runs { |
| 157 | if _, err := s.startLocked(r, cl); err != nil { |
| 158 | errs = append(errs, err) |
| 159 | } |
| 160 | } |
| 161 | return errors.Join(errs...) |
| 162 | } |
| 163 | |
| 164 | // Detach drops the current connection. Local listeners stay open (and refuse |
| 165 | // data until re-attached); remote listeners are closed. |
| 166 | func (s *Set) Detach() { |
| 167 | s.mu.Lock() |
| 168 | defer s.mu.Unlock() |
| 169 | s.ssh = nil |
| 170 | for _, r := range s.runs { |
| 171 | if r.remote != nil { |
| 172 | _ = r.remote.Close() |
| 173 | r.remote = nil |
| 174 | } |
| 175 | if r.up { |
| 176 | r.up = false |
| 177 | s.emit(Event{Spec: r.spec, Up: false}) |
| 178 | } |
| 179 | } |
| 180 | } |
| 181 | |
| 182 | // Close stops all forwards and releases every listener. |
| 183 | func (s *Set) Close() { |
| 184 | s.mu.Lock() |
| 185 | runs := s.runs |
| 186 | s.runs = map[string]*runner{} |
| 187 | s.ssh = nil |
| 188 | s.mu.Unlock() |
| 189 | for _, r := range runs { |
| 190 | s.stopRunner(r, true) |
| 191 | } |
| 192 | } |
| 193 | |
| 194 | // startLocked starts (or restarts) r on cl. Caller holds s.mu. |
| 195 | func (s *Set) startLocked(r *runner, cl *ssh.Client) (string, error) { |
| 196 | if r.spec.Direction == Local { |
| 197 | return s.startLocalLocked(r, cl) |
| 198 | } |
| 199 | return s.startRemoteLocked(r, cl) |
| 200 | } |
| 201 | |
| 202 | func (s *Set) startLocalLocked(r *runner, cl *ssh.Client) (string, error) { |
| 203 | if r.local == nil { |
| 204 | ln, err := net.Listen("tcp", r.spec.BindAddr) |
| 205 | if err != nil { |
| 206 | r.lastErr = wrapBind(err) |
| 207 | s.emit(Event{Spec: r.spec, Up: false, Err: r.lastErr}) |
| 208 | return "", r.lastErr |
| 209 | } |
| 210 | r.local = ln |
| 211 | r.acceptWG.Add(1) |
| 212 | go s.acceptLocal(r) |
| 213 | } |
| 214 | r.up = true |
| 215 | r.lastErr = nil |
| 216 | bound := r.local.Addr().String() |
| 217 | s.emit(Event{Spec: r.spec, Up: true, BoundAddr: bound}) |
| 218 | return bound, nil |
| 219 | } |
| 220 | |
| 221 | // acceptLocal accepts on the persistent local listener. Each accepted conn is |
| 222 | // forwarded through whatever ssh client is current at dial time; when detached |
| 223 | // (ssh == nil) the conn is refused. |
| 224 | func (s *Set) acceptLocal(r *runner) { |
| 225 | defer r.acceptWG.Done() |
| 226 | for { |
| 227 | conn, err := r.local.Accept() |
| 228 | if err != nil { |
| 229 | select { |
| 230 | case <-r.stop: |
| 231 | return |
| 232 | default: |
| 233 | return // listener closed |
| 234 | } |
| 235 | } |
| 236 | go s.handleLocalConn(r, conn) |
| 237 | } |
| 238 | } |
| 239 | |
| 240 | func (s *Set) handleLocalConn(r *runner, local net.Conn) { |
| 241 | s.mu.Lock() |
| 242 | cl := s.ssh |
| 243 | s.mu.Unlock() |
| 244 | if cl == nil { |
| 245 | _ = local.Close() |
| 246 | return |
| 247 | } |
| 248 | remote, err := cl.Dial("tcp", r.spec.TargetAddr) |
| 249 | if err != nil { |
| 250 | _ = local.Close() |
| 251 | return |
| 252 | } |
| 253 | pipe(local, remote) |
| 254 | } |
| 255 | |
| 256 | func (s *Set) startRemoteLocked(r *runner, cl *ssh.Client) (string, error) { |
| 257 | ln, err := cl.Listen("tcp", r.spec.BindAddr) |
| 258 | if err != nil { |
| 259 | r.lastErr = wrapBind(err) |
| 260 | s.emit(Event{Spec: r.spec, Up: false, Err: r.lastErr}) |
| 261 | return "", r.lastErr |
| 262 | } |
| 263 | r.remote = ln |
| 264 | r.up = true |
| 265 | r.lastErr = nil |
| 266 | go s.acceptRemote(r, ln) |
| 267 | bound := ln.Addr().String() |
| 268 | s.emit(Event{Spec: r.spec, Up: true, BoundAddr: bound}) |
| 269 | return bound, nil |
| 270 | } |
| 271 | |
| 272 | func (s *Set) acceptRemote(r *runner, ln net.Listener) { |
| 273 | defer func() { |
| 274 | s.mu.Lock() |
| 275 | // Only clear state for the still-current listener. Replace/Detach may |
| 276 | // have already swapped r.remote or marked the runner down. |
| 277 | if r.remote == ln { |
| 278 | r.remote = nil |
| 279 | if r.up { |
| 280 | r.up = false |
| 281 | s.emit(Event{Spec: r.spec, Up: false}) |
| 282 | } |
| 283 | } |
| 284 | s.mu.Unlock() |
| 285 | }() |
| 286 | for { |
| 287 | remote, err := ln.Accept() |
| 288 | if err != nil { |
| 289 | return |
| 290 | } |
| 291 | go func() { |
| 292 | local, derr := net.Dial("tcp", r.spec.TargetAddr) |
| 293 | if derr != nil { |
| 294 | _ = remote.Close() |
| 295 | return |
| 296 | } |
| 297 | pipe(remote, local) |
| 298 | }() |
| 299 | } |
| 300 | } |
| 301 | |
| 302 | func (s *Set) stopRunner(r *runner, closeLocal bool) { |
| 303 | close(r.stop) |
| 304 | |
| 305 | // acceptRemote also retires r.remote/r.up when Accept exits. Move the |
| 306 | // active listener and state out under the same lock so Remove/Replace/Close |
| 307 | // cannot race that deferred cleanup. |
| 308 | s.mu.Lock() |
| 309 | remote := r.remote |
| 310 | r.remote = nil |
| 311 | local := r.local |
| 312 | wasUp := r.up |
| 313 | r.up = false |
| 314 | s.mu.Unlock() |
| 315 | |
| 316 | if remote != nil { |
| 317 | _ = remote.Close() |
| 318 | } |
| 319 | if closeLocal && local != nil { |
| 320 | _ = local.Close() |
| 321 | } |
| 322 | r.acceptWG.Wait() |
| 323 | if closeLocal { |
| 324 | s.mu.Lock() |
| 325 | if r.local == local { |
| 326 | r.local = nil |
| 327 | } |
| 328 | s.mu.Unlock() |
| 329 | } |
| 330 | if wasUp { |
| 331 | s.emit(Event{Spec: r.spec, Up: false}) |
| 332 | } |
| 333 | } |
| 334 | |
| 335 | func (s *Set) emit(e Event) { |
| 336 | if s.onEvent != nil { |
| 337 | s.onEvent(e) |
| 338 | } |
| 339 | } |
| 340 | |
| 341 | // pipe copies bidirectionally between a and b, closing both when either side |
| 342 | // ends. Half-close is best-effort via CloseWrite when supported. |
| 343 | func pipe(a, b net.Conn) { |
| 344 | done := make(chan struct{}, 2) |
| 345 | cp := func(dst, src net.Conn) { |
| 346 | _, _ = io.Copy(dst, src) |
| 347 | if cw, ok := dst.(interface{ CloseWrite() error }); ok { |
| 348 | _ = cw.CloseWrite() |
| 349 | } |
| 350 | done <- struct{}{} |
| 351 | } |
| 352 | go cp(a, b) |
| 353 | go cp(b, a) |
| 354 | <-done |
| 355 | _ = a.Close() |
| 356 | _ = b.Close() |
| 357 | } |
| 358 | |
| 359 | func wrapBind(err error) error { |
| 360 | if isAddrInUse(err) { |
| 361 | return errors.Join(ErrBindBusy, err) |
| 362 | } |
| 363 | return err |
| 364 | } |
| 365 |