Simplified how values are returned from the GenerateSafePrime

We can use `defer waitGroup.Wait()` to make sure it's called regardless
which path we go. This we, we can call `return ...` directly from all
select cases and we don't need named parameters.
This commit is contained in:
Piotr Dyraga committed 2018-07-30 16:04:00 +02:00
1 parent 82f050d7b2
commit 1200243127
1 file changed
+7 -9
+7 -9
View File
@@ -63,7 +63,7 @@ func GenerateSafePrime(
concurrencyLevel int, concurrencyLevel int,
timeout time.Duration, timeout time.Duration,
random io.Reader, random io.Reader,
) (p *big.Int, q *big.Int, err error) { ) (*big.Int, *big.Int, error) {
if bitLen < 6 { if bitLen < 6 {
return nil, nil, errors.New("safe prime size must be at least 6 bits") return nil, nil, errors.New("safe prime size must be at least 6 bits")
} }
@@ -71,10 +71,11 @@ func GenerateSafePrime(
primeChan := make(chan safePrime, concurrencyLevel) primeChan := make(chan safePrime, concurrencyLevel)
errChan := make(chan error, concurrencyLevel) errChan := make(chan error, concurrencyLevel)
waitGroup := &sync.WaitGroup{}
defer close(primeChan) defer close(primeChan)
defer close(errChan) defer close(errChan)
defer waitGroup.Wait()
waitGroup := &sync.WaitGroup{}
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
@@ -94,16 +95,13 @@ func GenerateSafePrime(
select { select {
case result := <-primeChan: case result := <-primeChan:
cancel() cancel()
p, q, err = result.p, result.q, nil return result.p, result.q, nil
case err := <-errChan: case err := <-errChan:
cancel() cancel()
p, q, err = nil, nil, err return nil, nil, err
case <-ctx.Done(): case <-ctx.Done():
p, q, err = nil, nil, fmt.Errorf("generator timed out after %v", timeout) return nil, nil, fmt.Errorf("generator timed out after %v", timeout)
} }
waitGroup.Wait()
return
} }
type safePrime struct { type safePrime struct {