-
Notifications
You must be signed in to change notification settings - Fork 35
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
readerutil: add ConnReadN function (#275)
* readutil: add ConnReadN function * misc improvements * revert #273 with explaination * introduce: NContext * misc changes * use context / set timeout * bump go -> 1.21 * hide/accept timeout error if we get any data * fix accepted error timeout logic * fix race condition * improved explaination of testcase * include pollerr in IsTimeout()
- Loading branch information
1 parent
7240002
commit 62487ea
Showing
13 changed files
with
413 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -15,7 +15,7 @@ jobs: | |
- name: Set up Go | ||
uses: actions/setup-go@v4 | ||
with: | ||
go-version: 1.20.x | ||
go-version: 1.21.x | ||
- name: Run golangci-lint | ||
uses: golangci/[email protected] | ||
with: | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,91 @@ | ||
package contextutil | ||
|
||
import "context" | ||
|
||
// A problematic situation when implementing context in a function | ||
// is when that function has more than one return values | ||
// if function has only one return value we can safely wrap it something like this | ||
/* | ||
func DoSomething() error {} | ||
ch := make(chan error) | ||
go func() { | ||
ch <- DoSomething() | ||
}() | ||
select { | ||
case err := <-ch: | ||
// handle error | ||
case <-ctx.Done(): | ||
// handle context cancelation | ||
} | ||
*/ | ||
// but what if we have more than one value to return? | ||
// we can use generics and a struct and that is what we are doing here | ||
// here we use struct and generics to store return values of a function | ||
// instead of storing it in a []interface{} | ||
|
||
type twoValueCtx[T1 any, T2 any] struct { | ||
var1 T1 | ||
var2 T2 | ||
} | ||
|
||
type threeValueCtx[T1 any, T2 any, T3 any] struct { | ||
var1 T1 | ||
var2 T2 | ||
var3 T3 | ||
} | ||
|
||
// ExecFunc implements context for a function which has no return values | ||
// and executes that function. if context is cancelled before function returns | ||
// it will return context error otherwise it will return nil | ||
func ExecFunc(ctx context.Context, fn func()) error { | ||
ch := make(chan struct{}) | ||
go func() { | ||
fn() | ||
ch <- struct{}{} | ||
}() | ||
select { | ||
case <-ch: | ||
return nil | ||
case <-ctx.Done(): | ||
return ctx.Err() | ||
} | ||
} | ||
|
||
// ExecFuncWithTwoReturns wraps a function which has two return values given that last one is error | ||
// and executes that function in a goroutine there by implementing context | ||
// if context is cancelled before function returns it will return context error | ||
// otherwise it will return function's return values | ||
func ExecFuncWithTwoReturns[T1 any](ctx context.Context, fn func() (T1, error)) (T1, error) { | ||
ch := make(chan twoValueCtx[T1, error]) | ||
go func() { | ||
x, y := fn() | ||
ch <- twoValueCtx[T1, error]{var1: x, var2: y} | ||
}() | ||
select { | ||
case <-ctx.Done(): | ||
var tmp T1 | ||
return tmp, ctx.Err() | ||
case v := <-ch: | ||
return v.var1, v.var2 | ||
} | ||
} | ||
|
||
// ExecFuncWithThreeReturns wraps a function which has three return values given that last one is error | ||
// and executes that function in a goroutine there by implementing context | ||
// if context is cancelled before function returns it will return context error | ||
// otherwise it will return function's return values | ||
func ExecFuncWithThreeReturns[T1 any, T2 any](ctx context.Context, fn func() (T1, T2, error)) (T1, T2, error) { | ||
ch := make(chan threeValueCtx[T1, T2, error]) | ||
go func() { | ||
x, y, z := fn() | ||
ch <- threeValueCtx[T1, T2, error]{var1: x, var2: y, var3: z} | ||
}() | ||
select { | ||
case <-ctx.Done(): | ||
var tmp1 T1 | ||
var tmp2 T2 | ||
return tmp1, tmp2, ctx.Err() | ||
case v := <-ch: | ||
return v.var1, v.var2, v.var3 | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,110 @@ | ||
package contextutil_test | ||
|
||
import ( | ||
"context" | ||
"errors" | ||
"testing" | ||
"time" | ||
|
||
contextutil "github.com/projectdiscovery/utils/context" | ||
) | ||
|
||
func TestExecFuncWithTwoReturns(t *testing.T) { | ||
t.Run("function completes before context cancellation", func(t *testing.T) { | ||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | ||
defer cancel() | ||
|
||
fn := func() (int, error) { | ||
time.Sleep(1 * time.Second) | ||
return 42, nil | ||
} | ||
|
||
val, err := contextutil.ExecFuncWithTwoReturns(ctx, fn) | ||
if err != nil { | ||
t.Errorf("Unexpected error: %v", err) | ||
} | ||
if val != 42 { | ||
t.Errorf("Unexpected return value: got %v, want 42", val) | ||
} | ||
}) | ||
|
||
t.Run("context cancelled before function completes", func(t *testing.T) { | ||
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) | ||
defer cancel() | ||
|
||
fn := func() (int, error) { | ||
time.Sleep(2 * time.Second) | ||
return 42, nil | ||
} | ||
|
||
_, err := contextutil.ExecFuncWithTwoReturns(ctx, fn) | ||
if !errors.Is(err, context.DeadlineExceeded) { | ||
t.Errorf("Expected context deadline exceeded error, got: %v", err) | ||
} | ||
}) | ||
} | ||
|
||
func TestExecFuncWithThreeReturns(t *testing.T) { | ||
t.Run("function completes before context cancellation", func(t *testing.T) { | ||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | ||
defer cancel() | ||
|
||
fn := func() (int, string, error) { | ||
time.Sleep(1 * time.Second) | ||
return 42, "hello", nil | ||
} | ||
|
||
val1, val2, err := contextutil.ExecFuncWithThreeReturns(ctx, fn) | ||
if err != nil { | ||
t.Errorf("Unexpected error: %v", err) | ||
} | ||
if val1 != 42 || val2 != "hello" { | ||
t.Errorf("Unexpected return values: got %v and %v, want 42 and 'hello'", val1, val2) | ||
} | ||
}) | ||
|
||
t.Run("context cancelled before function completes", func(t *testing.T) { | ||
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) | ||
defer cancel() | ||
|
||
fn := func() (int, string, error) { | ||
time.Sleep(2 * time.Second) | ||
return 42, "hello", nil | ||
} | ||
|
||
_, _, err := contextutil.ExecFuncWithThreeReturns(ctx, fn) | ||
if !errors.Is(err, context.DeadlineExceeded) { | ||
t.Errorf("Expected context deadline exceeded error, got: %v", err) | ||
} | ||
}) | ||
} | ||
|
||
func TestExecFunc(t *testing.T) { | ||
t.Run("function completes before context cancellation", func(t *testing.T) { | ||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) | ||
defer cancel() | ||
|
||
fn := func() { | ||
time.Sleep(1 * time.Second) | ||
} | ||
|
||
err := contextutil.ExecFunc(ctx, fn) | ||
if err != nil { | ||
t.Errorf("Unexpected error: %v", err) | ||
} | ||
}) | ||
|
||
t.Run("context cancelled before function completes", func(t *testing.T) { | ||
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) | ||
defer cancel() | ||
|
||
fn := func() { | ||
time.Sleep(2 * time.Second) | ||
} | ||
|
||
err := contextutil.ExecFunc(ctx, fn) | ||
if err != context.DeadlineExceeded { | ||
t.Errorf("Expected context deadline exceeded error, got: %v", err) | ||
} | ||
}) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,6 +1,6 @@ | ||
module github.com/projectdiscovery/utils | ||
|
||
go 1.20 | ||
go 1.21 | ||
|
||
require ( | ||
github.com/Masterminds/semver/v3 v3.2.1 | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.