-
Notifications
You must be signed in to change notification settings - Fork 10
/
stream_handler.go
54 lines (47 loc) · 1.45 KB
/
stream_handler.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
package stream_echo
import (
"context"
"github.com/aperturerobotics/bifrost/link"
"github.com/aperturerobotics/controllerbus/bus"
"github.com/aperturerobotics/util/ioproxy"
"github.com/sirupsen/logrus"
)
// MountedStreamHandler implements the mounted stream handler.
type MountedStreamHandler struct {
// le is the logger entry
le *logrus.Entry
// bus is the controller bus
bus bus.Bus
}
// NewMountedStreamHandler constructs the mounted stream handler.
func NewMountedStreamHandler(le *logrus.Entry, bus bus.Bus) (*MountedStreamHandler, error) {
return &MountedStreamHandler{le: le, bus: bus}, nil
}
// HandleMountedStream handles an incoming mounted stream.
// Any returned error indicates the stream should be closed.
// This function should return as soon as possible, and start
// additional goroutines to manage the lifecycle of the stream.
func (m *MountedStreamHandler) HandleMountedStream(
ctx context.Context,
strm link.MountedStream,
) error {
_, elRef, err := m.bus.AddDirective(
link.NewEstablishLinkWithPeer(strm.GetLink().GetLocalPeer(), strm.GetPeerID()),
nil,
)
if err != nil {
return err
}
go func() {
defer elRef.Release()
m.le.Debug("echoing stream")
s := strm.GetStream()
subCtx, subCtxCancel := context.WithCancel(ctx)
defer subCtxCancel()
ioproxy.ProxyStreams(s, s, subCtxCancel)
// wait to release EstablishLink ref
<-subCtx.Done()
}()
return nil
}
var _ link.MountedStreamHandler = ((*MountedStreamHandler)(nil))