-
Notifications
You must be signed in to change notification settings - Fork 175
/
corruptedConnection.go
97 lines (80 loc) · 2.82 KB
/
corruptedConnection.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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
package attacknetwork
import (
"context"
"errors"
"fmt"
"io"
"github.com/rs/zerolog"
"github.com/onflow/flow-go/insecure"
"github.com/onflow/flow-go/module/component"
"github.com/onflow/flow-go/module/irrecoverable"
)
// CorruptedNodeConnection abstracts connection between an attack orchestrator to a corruptible conduit factory (ccf)
// through the attack network.
type CorruptedNodeConnection struct {
component.Component
cm *component.ComponentManager
logger zerolog.Logger
inboundHandler func(*insecure.Message) // handler for incoming messages from corruptible conduit factories.
outbound insecure.CorruptibleConduitFactory_ProcessAttackerMessageClient // from orchestrator to ccf.
inbound insecure.CorruptibleConduitFactory_ConnectAttackerClient // from ccf to orchestrator.
}
func NewCorruptedNodeConnection(
logger zerolog.Logger,
inboundHandler func(message *insecure.Message),
outbound insecure.CorruptibleConduitFactory_ProcessAttackerMessageClient,
inbound insecure.CorruptibleConduitFactory_ConnectAttackerClient) *CorruptedNodeConnection {
c := &CorruptedNodeConnection{
logger: logger.With().Str("component", "corrupted-connector").Logger(),
outbound: outbound,
inbound: inbound,
inboundHandler: inboundHandler,
}
cm := component.NewComponentManagerBuilder().
AddWorker(func(ctx irrecoverable.SignalerContext, ready component.ReadyFunc) {
ready()
c.receiveLoop()
<-ctx.Done()
}).Build()
c.Component = cm
c.cm = cm
return c
}
// SendMessage sends the message from orchestrator to the corrupted conduit factory.
func (c *CorruptedNodeConnection) SendMessage(message *insecure.Message) error {
err := c.outbound.Send(message)
if err != nil {
return fmt.Errorf("could not send message: %w", err)
}
return nil
}
// receiveLoop implements the continuous procedure of reading from inbound stream of this connection, which
// is established from the remote ccf to the local attack orchestrator.
func (c *CorruptedNodeConnection) receiveLoop() {
c.logger.Info().Msg("receive loop started")
for {
select {
case <-c.cm.ShutdownSignal():
// connection closed
c.logger.Info().Msg("receive loop terminated")
return
default:
msg, err := c.inbound.Recv()
if err == io.EOF || errors.Is(c.inbound.Context().Err(), context.Canceled) {
c.logger.Warn().Msg("inbound stream closed")
return
} else if err != nil {
c.logger.Error().Err(err).Msg("error reading inbound stream")
}
c.inboundHandler(msg)
}
}
}
// CloseConnection closes the connection to the corrupted conduit factory.
func (c *CorruptedNodeConnection) CloseConnection() error {
err := c.outbound.CloseSend()
if err != nil {
return fmt.Errorf("could not close outbound connection: %w", err)
}
return nil
}