/
tcp.go
97 lines (86 loc) · 2.26 KB
/
tcp.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 gohc
import (
"crypto/tls"
"fmt"
"io"
"net"
"time"
)
// TcpOpt Describes the TCP health check specific options.
type TcpOpt struct {
// TCP specific payload.
// Empty payloads imply a connect-only health check.
Send *Payload
// When checking the response, “fuzzy” matching is performed such that each
// binary block must be found, and in the order specified, but not
// necessarily contiguous.
Receive []*Payload
// Timeout for connection and for each receive data. If left empty (default to 5s)
Timeout time.Duration
// TlsEnabled set to true if the gRPC health check request should be sent over TLS.
TlsEnabled bool
// TlsConfig specifies the TLS configuration to use for TLS enabled gRPC health check requests.
TlsConfig *tls.Config
// AltPort specifies the port to use for gRPC health check requests.
// If left empty it taks the port from host during check.
AltPort uint32
}
type TcpHealthCheck struct {
opt *TcpOpt
}
func NewTcpHealthCheck(opt *TcpOpt) *TcpHealthCheck {
return &TcpHealthCheck{
opt: opt,
}
}
func (h *TcpHealthCheck) Check(host string) error {
netConn, err := h.makeNetConn(host)
if err != nil {
return err
}
defer netConn.Close()
timeout := h.opt.Timeout
if timeout == 0 {
timeout = 5 * time.Second
}
if h.opt.Send != nil {
_, err = netConn.Write(h.opt.Send.GetData())
if err != nil {
return err
}
}
for _, toReceive := range h.opt.Receive {
err := netConn.SetReadDeadline(time.Now().Add(timeout))
if err != nil {
return err
}
buf := make([]byte, len(toReceive.GetData()))
n, err := io.ReadFull(netConn, buf)
if err != nil {
return fmt.Errorf("failed to read %d bytes: %v", len(toReceive.GetData()), err)
}
got := buf[0:n]
if string(got) != string(toReceive.GetData()) {
return fmt.Errorf("expected %s, got %s", string(toReceive.GetData()), string(got))
}
}
return nil
}
func (h *TcpHealthCheck) makeNetConn(host string) (net.Conn, error) {
var err error
timeout := h.opt.Timeout
if timeout == 0 {
timeout = 5 * time.Second
}
host, err = FormatHost(host, h.opt.AltPort)
if err != nil {
return nil, err
}
dialer := &net.Dialer{
Timeout: timeout,
}
if h.opt.TlsEnabled {
return tls.DialWithDialer(dialer, "tcp", host, h.opt.TlsConfig)
}
return dialer.Dial("tcp", host)
}