From 867a243fef2743ff50d51393d38b1874a3f1a8cc Mon Sep 17 00:00:00 2001
From: Simone Basso <bassosimone@gmail.com>
Date: Sun, 28 Aug 2022 12:17:31 +0200
Subject: [PATCH] refactor(oohelperd): make performing additional measurements
 easier (#889)

This diff refactors oohelperd to make performing additional measurements easier. We need:

1. to run the DNS task _before_ other tasks such that we can measure both IP addresses returned by the TH and the ones returned by the probe. When we'll introduce TLS measurements, this will allow us to validate probe-provided IP addresses inside the TH call. If probe-provided addresses work with TLS, they are legitimate for the domain.

2. to tie the number of TCP measurements to a list of endpoints collected by the probe _or_ the TH rather than just to the one provided by the probe. Anticipating this change, let us refactor how we read the results of the TCP task to make it independent of the number of addresses provided by the probe.

This work is part of https://github.com/ooni/probe/issues/2237
---
 internal/cmd/oohelperd/measure.go | 42 +++++++++++++++++++------------
 1 file changed, 26 insertions(+), 16 deletions(-)

diff --git a/internal/cmd/oohelperd/measure.go b/internal/cmd/oohelperd/measure.go
index 5035642..887d2a0 100644
--- a/internal/cmd/oohelperd/measure.go
+++ b/internal/cmd/oohelperd/measure.go
@@ -43,6 +43,22 @@ func measure(ctx context.Context, config *handler, creq *ctrlRequest) (*ctrlResp
 		})
 	}
 
+	// wait for DNS measurements to complete
+	wg.Wait()
+
+	// start assembling the response
+	cresp := new(ctrlResponse)
+	select {
+	case cresp.DNS = <-dnsch:
+	default:
+		// we need to emit a non-nil Addrs to match exactly
+		// the behavior of the legacy TH
+		cresp.DNS = ctrlDNSResult{
+			Failure: nil,
+			Addrs:   []string{},
+		}
+	}
+
 	// tcpconnect: start
 	tcpconnch := make(chan tcpResultPair, len(creq.TCPConnect))
 	for _, endpoint := range creq.TCPConnect {
@@ -67,26 +83,20 @@ func measure(ctx context.Context, config *handler, creq *ctrlRequest) (*ctrlResp
 		Wg:                wg,
 	})
 
-	// wait for measurement steps to complete
+	// wait for endpoint measurements to complete
 	wg.Wait()
 
-	// assemble response
-	cresp := new(ctrlResponse)
-	select {
-	case cresp.DNS = <-dnsch:
-	default:
-		// we need to emit a non-nil Addrs to match exactly
-		// the behavior of the legacy TH
-		cresp.DNS = ctrlDNSResult{
-			Failure: nil,
-			Addrs:   []string{},
-		}
-	}
+	// continue assembling the response
 	cresp.HTTPRequest = <-httpch
 	cresp.TCPConnect = make(map[string]ctrlTCPResult)
-	for len(cresp.TCPConnect) < len(creq.TCPConnect) {
-		tcpconn := <-tcpconnch
-		cresp.TCPConnect[tcpconn.Endpoint] = tcpconn.Result
+Loop:
+	for {
+		select {
+		case tcpconn := <-tcpconnch:
+			cresp.TCPConnect[tcpconn.Endpoint] = tcpconn.Result
+		default:
+			break Loop
+		}
 	}
 
 	return cresp, nil