/
tcputil.go
133 lines (114 loc) · 3.39 KB
/
tcputil.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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
// Copyright 2015 The Vanadium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// package tcputil contains functions commonly used to manipulate TCP
// connections.
package tcputil
import (
"net"
"time"
"v.io/v23/context"
"v.io/v23/flow"
"v.io/x/ref/runtime/protocols/lib/framer"
)
const keepAlivePeriod = 30 * time.Second
// EnableTCPKeepAlive enabled the KeepAlive option on a TCP connection.
//
// Some cloud providers (like Google Compute Engine) blackhole inactive TCP
// connections, we need to set TCP keep alive option to prevent that.
// See: https://developers.google.com/compute/docs/troubleshooting#communicatewithinternet
//
// The same problem can happen when one end of a TCP connection dies and the
// TCP FIN or RST packet doesn't reach the other end, e.g. when the machine
// dies, falls off the network, or when there is packet loss. So, it is best to
// enable this option for all TCP connections.
func EnableTCPKeepAlive(conn net.Conn) error {
if tcpconn, ok := conn.(*net.TCPConn); ok {
if err := tcpconn.SetKeepAlivePeriod(keepAlivePeriod); err != nil {
return err
}
return tcpconn.SetKeepAlive(true)
}
return nil
}
type TCP struct{}
// Dial dials a net.Conn to a the specific address and adds framing to the connection.
func (TCP) Dial(ctx *context.T, network, address string, timeout time.Duration) (flow.Conn, error) {
conn, err := net.DialTimeout(network, address, timeout)
if err != nil {
return nil, err
}
if err := EnableTCPKeepAlive(conn); err != nil {
return nil, err
}
return NewTCPConn(conn), nil
}
// Resolve performs a DNS resolution on the provided network and address.
func (TCP) Resolve(ctx *context.T, network, address string) (string, []string, error) {
addrs, err := TCPResolveAddrs(ctx, address)
return network, addrs, err
}
// Listen returns a listener that sets KeepAlive on all accepted connections.
// Connections returned from the listener will be framed.
func (TCP) Listen(ctx *context.T, network, address string) (flow.Listener, error) {
ln, err := net.Listen(network, address)
if err != nil {
return nil, err
}
return &tcpListener{ln}, nil
}
// tcpListener is a wrapper around net.Listener that sets KeepAlive on all
// accepted connections and returns framed flow.Conns.
type tcpListener struct {
netLn net.Listener
}
func (ln *tcpListener) Accept(ctx *context.T) (flow.Conn, error) {
conn, err := ln.netLn.Accept()
if err != nil {
return nil, err
}
if err := EnableTCPKeepAlive(conn); err != nil {
return nil, err
}
return NewTCPConn(conn), nil
}
func (ln *tcpListener) Addr() net.Addr {
return ln.netLn.Addr()
}
func (ln *tcpListener) Close() error {
return ln.netLn.Close()
}
func NewTCPConn(c net.Conn) flow.Conn {
return tcpConn{
framer.New(c),
c.LocalAddr(),
c.RemoteAddr(),
}
}
type tcpConn struct {
flow.MsgReadWriteCloser
localAddr net.Addr
remoteAddr net.Addr
}
func (c tcpConn) LocalAddr() net.Addr {
return c.localAddr
}
func (c tcpConn) RemoteAddr() net.Addr {
return c.remoteAddr
}
func TCPResolveAddrs(ctx *context.T, address string) ([]string, error) {
host, port, err := net.SplitHostPort(address)
if err != nil {
return nil, err
}
var addrs []string
ips, err := net.LookupIP(host)
if err != nil {
return nil, err
}
s := ":" + port
for _, ip := range ips {
addrs = append(addrs, ip.String()+s)
}
return addrs, nil
}