Skip to content

Commit c0b7fc9

Browse files
firefartfirefart
committed
fix #47 
1 parent e709bda commit c0b7fc9

File tree

4 files changed

+59
-159
lines changed

4 files changed

+59
-159
lines changed

‎internal/helper/connection.go

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package helper
33
import (
44
"context"
55
"errors"
6+
"fmt"
67
"io"
78
"net"
89
"time"
@@ -17,13 +18,18 @@ func ConnectionRead(ctx context.Context, conn net.Conn, timeout time.Duration) (
1718
ctx, cancel := context.WithTimeout(ctx, timeout)
1819
defer cancel()
1920

21+
// need this otherwise the read call is blocking forever
22+
if err := conn.SetReadDeadline(time.Now().Add(timeout)); err != nil {
23+
return nil, fmt.Errorf("could not set read deadline: %v", err)
24+
}
25+
2026
bufLen := 1024
27+
buf := make([]byte, bufLen)
2128
for {
2229
select {
2330
case <-ctx.Done():
2431
return nil, ctx.Err()
2532
default:
26-
buf := make([]byte, bufLen)
2733
i, err := conn.Read(buf)
2834
if err != nil {
2935
if err != io.EOF {
@@ -53,6 +59,11 @@ func ConnectionWrite(ctx context.Context, conn net.Conn, data []byte, timeout ti
5359
ctx, cancel := context.WithTimeout(ctx, timeout)
5460
defer cancel()
5561

62+
// need this otherwise the read call is blocking forever
63+
if err := conn.SetWriteDeadline(time.Now().Add(timeout)); err != nil {
64+
return fmt.Errorf("could not set write deadline: %v", err)
65+
}
66+
5667
for {
5768
select {
5869
case <-ctx.Done():

‎internal/socksimplementations/socksturntcphandler.go

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,18 +117,42 @@ func (s *SocksTurnTCPHandler) Refresh(ctx context.Context) {
117117

118118
const bufferLength = 1024 * 100
119119

120+
type readDeadline interface {
121+
SetReadDeadline(time.Time) error
122+
}
123+
type writeDeadline interface {
124+
SetWriteDeadline(time.Time) error
125+
}
126+
120127
// ReadFromClient is used to copy data
121128
func (s *SocksTurnTCPHandler) ReadFromClient(ctx context.Context, client io.ReadCloser, remote io.WriteCloser) error {
122129
for {
123130
// anonymous func for defer
124131
// this might not be the fastest, but it does the trick
132+
// in this case the timeout is per buffer read/write to support long
133+
// running downloads.
125134
err := func() error {
126-
ctx, cancel := context.WithTimeout(ctx, s.Timeout)
135+
timeOut := time.Now().Add(s.Timeout)
136+
137+
ctx, cancel := context.WithDeadline(ctx, timeOut)
127138
defer cancel()
139+
128140
select {
129141
case <-ctx.Done():
130142
return ctx.Err()
131143
default:
144+
if c, ok := remote.(writeDeadline); ok {
145+
if err := c.SetWriteDeadline(timeOut); err != nil {
146+
return fmt.Errorf("could not set write deadline on remote: %v", err)
147+
}
148+
}
149+
150+
if c, ok := client.(readDeadline); ok {
151+
if err := c.SetReadDeadline(timeOut); err != nil {
152+
return fmt.Errorf("could not set read deadline on client: %v", err)
153+
}
154+
}
155+
132156
i, err := io.CopyN(remote, client, bufferLength)
133157
if errors.Is(err, io.EOF) {
134158
return nil
@@ -150,13 +174,30 @@ func (s *SocksTurnTCPHandler) ReadFromRemote(ctx context.Context, remote io.Read
150174
for {
151175
// anonymous func for defer
152176
// this might not be the fastest, but it does the trick
177+
// in this case the timeout is per buffer read/write to support long
178+
// running downloads.
153179
err := func() error {
154-
ctx, cancel := context.WithTimeout(ctx, s.Timeout)
180+
timeOut := time.Now().Add(s.Timeout)
181+
182+
ctx, cancel := context.WithDeadline(ctx, timeOut)
155183
defer cancel()
184+
156185
select {
157186
case <-ctx.Done():
158187
return ctx.Err()
159188
default:
189+
if c, ok := client.(writeDeadline); ok {
190+
if err := c.SetWriteDeadline(timeOut); err != nil {
191+
return fmt.Errorf("could not set write deadline on client: %v", err)
192+
}
193+
}
194+
195+
if c, ok := remote.(readDeadline); ok {
196+
if err := c.SetReadDeadline(timeOut); err != nil {
197+
return fmt.Errorf("could not set read deadline on remote: %v", err)
198+
}
199+
}
200+
160201
i, err := io.CopyN(client, remote, bufferLength)
161202
if errors.Is(err, io.EOF) {
162203
return nil

‎internal/socksimplementations/socksturnudphandler.go

Lines changed: 0 additions & 152 deletions
This file was deleted.

‎main.go

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ func main() {
5252
&cli.StringFlag{Name: "turnserver", Aliases: []string{"s"}, Required: true, Usage: "turn server to connect to in the format host:port"},
5353
&cli.BoolFlag{Name: "tls", Value: false, Usage: "Use TLS/DTLS on connecting to the STUN or TURN server"},
5454
&cli.StringFlag{Name: "protocol", Value: "udp", Usage: "protocol to use when connecting to the TURN server. Supported values: tcp and udp"},
55-
&cli.DurationFlag{Name: "timeout", Value: 5 * time.Second, Usage: "connect timeout to turn server"},
55+
&cli.DurationFlag{Name: "timeout", Value: 2 * time.Second, Usage: "connect timeout to turn server"},
5656
},
5757
Before: func(ctx *cli.Context) error {
5858
if ctx.Bool("debug") {
@@ -86,7 +86,7 @@ func main() {
8686
&cli.StringFlag{Name: "turnserver", Aliases: []string{"s"}, Required: true, Usage: "turn server to connect to in the format host:port"},
8787
&cli.BoolFlag{Name: "tls", Value: false, Usage: "Use TLS/DTLS on connecting to the STUN or TURN server"},
8888
&cli.StringFlag{Name: "protocol", Value: "udp", Usage: "protocol to use when connecting to the TURN server. Supported values: tcp and udp"},
89-
&cli.DurationFlag{Name: "timeout", Value: 5 * time.Second, Usage: "connect timeout to turn server"},
89+
&cli.DurationFlag{Name: "timeout", Value: 2 * time.Second, Usage: "connect timeout to turn server"},
9090
&cli.StringFlag{Name: "username", Aliases: []string{"u"}, Required: true, Usage: "username for the turn server"},
9191
&cli.StringFlag{Name: "password", Aliases: []string{"p"}, Required: true, Usage: "password for the turn server"},
9292
},
@@ -168,7 +168,7 @@ func main() {
168168
&cli.StringFlag{Name: "turnserver", Aliases: []string{"s"}, Required: true, Usage: "turn server to connect to in the format host:port"},
169169
&cli.BoolFlag{Name: "tls", Value: false, Usage: "Use TLS/DTLS on connecting to the STUN or TURN server"},
170170
&cli.StringFlag{Name: "protocol", Value: "udp", Usage: "protocol to use when connecting to the TURN server. Supported values: tcp and udp"},
171-
&cli.DurationFlag{Name: "timeout", Value: 5 * time.Second, Usage: "connect timeout to turn server"},
171+
&cli.DurationFlag{Name: "timeout", Value: 2 * time.Second, Usage: "connect timeout to turn server"},
172172
&cli.StringFlag{Name: "username", Aliases: []string{"u"}, Required: true, Usage: "username for the turn server"},
173173
&cli.StringFlag{Name: "password", Aliases: []string{"p"}, Required: true, Usage: "password for the turn server"},
174174
&cli.StringFlag{Name: "target", Aliases: []string{"t"}, Required: true, Usage: "Target to leak memory to in the form host:port. Should be a public server under your control"},
@@ -231,7 +231,7 @@ func main() {
231231
&cli.StringFlag{Name: "turnserver", Aliases: []string{"s"}, Required: true, Usage: "turn server to connect to in the format host:port"},
232232
&cli.BoolFlag{Name: "tls", Value: false, Usage: "Use TLS/DTLS on connecting to the STUN or TURN server"},
233233
&cli.StringFlag{Name: "protocol", Value: "udp", Usage: "protocol to use when connecting to the TURN server. Supported values: tcp and udp"},
234-
&cli.DurationFlag{Name: "timeout", Value: 5 * time.Second, Usage: "connect timeout to turn server"},
234+
&cli.DurationFlag{Name: "timeout", Value: 2 * time.Second, Usage: "connect timeout to turn server"},
235235
&cli.StringFlag{Name: "username", Aliases: []string{"u"}, Required: true, Usage: "username for the turn server"},
236236
&cli.StringFlag{Name: "password", Aliases: []string{"p"}, Required: true, Usage: "password for the turn server"},
237237
},

0 commit comments

Comments
 (0)