tor: add method Reconnect to reset connection

A new method, Reconnect, is added to tor controller which can be used to
reset the current connection. This will be later used in healthcheck to
help us reset the connection to Tor Daemon.

tor/controller.go

@@ -69,6 +69,14 @@ var (
 	// errCodeNotMatch is used when an expected response code is not
 	// returned.
 	errCodeNotMatch = errors.New("unexpected code")
+
+	// errTCNotStarted is used when we require the tor controller to be
+	// started while it's not.
+	errTCNotStarted = errors.New("tor controller must be started")
+
+	// errTCNotStarted is used when we require the tor controller to be
+	// not stopped while it is.
+	errTCStopped = errors.New("tor controller must not be stopped")
 )
 
 // Controller is an implementation of the Tor Control protocol. This is used in
@@ -169,6 +177,60 @@ func (c *Controller) Stop() error {
 	return c.conn.Close()
 }
 
+// Reconnect makes a new socket connection between the tor controller and
+// daemon. It will attempt to close the old connection, make a new connection
+// and authenticate, and finally reset the activeServiceID that the controller
+// is aware of.
+//
+// NOTE: Any old onion services will be removed once this function is called.
+// In the case of a Tor daemon restart, previously created onion services will
+// no longer be there. If the function is called without a Tor daemon restart,
+// because the control connection is reset, all the onion services belonging to
+// the old connection will be removed.
+func (c *Controller) Reconnect() error {
+	// Require the tor controller to be running when we want to reconnect.
+	// This means the started flag must be 1 and the stopped flag must be
+	// 0.
+	if c.started != 1 {
+		return errTCNotStarted
+	}
+	if c.stopped != 0 {
+		return errTCStopped
+	}
+
+	log.Info("Re-connectting tor controller")
+
+	// If we have an old connection, try to close it. We might receive an
+	// error if the connection has already been closed by Tor daemon(ie,
+	// daemon restarted), so we ignore the error here.
+	if c.conn != nil {
+		if err := c.conn.Close(); err != nil {
+			log.Debugf("closing old conn got err: %v", err)
+		}
+	}
+
+	// Make a new connection and authenticate.
+	conn, err := textproto.Dial("tcp", c.controlAddr)
+	if err != nil {
+		return fmt.Errorf("unable to connect to Tor server: %w", err)
+	}
+
+	c.conn = conn
+
+	// Authenticate the connection between the controller and Tor daemon.
+	if err := c.authenticate(); err != nil {
+		return err
+	}
+
+	// Reset the activeServiceID. This value would only be set if a
+	// previous onion service was created. Because the old connection has
+	// been closed at this point, the old onion service is no longer
+	// active.
+	c.activeServiceID = ""
+
+	return nil
+}
+
 // sendCommand sends a command to the Tor server and returns its response, as a
 // single space-delimited string, and code.
 func (c *Controller) sendCommand(command string) (int, string, error) {

tor/controller_test.go

@@ -276,3 +276,87 @@ func TestReadResponse(t *testing.T) {
 		})
 	}
 }
+
+// TestReconnectTCMustBeRunning checks that the tor controller must be running
+// while calling Reconnect.
+func TestReconnectTCMustBeRunning(t *testing.T) {
+	// Create a dummy controller.
+	c := &Controller{}
+
+	// Reconnect should fail because the TC is not started.
+	require.Equal(t, errTCNotStarted, c.Reconnect())
+
+	// Set the started flag.
+	c.started = 1
+
+	// Set the stopped flag so the TC is stopped.
+	c.stopped = 1
+
+	// Reconnect should fail because the TC is stopped.
+	require.Equal(t, errTCStopped, c.Reconnect())
+}
+
+// TestReconnectSucceed tests a reconnection will succeed when the tor
+// controller is up and running.
+func TestReconnectSucceed(t *testing.T) {
+	// Create mock server and client connection.
+	proxy := createTestProxy(t)
+	defer proxy.cleanUp()
+
+	// Create a tor controller and mark the controller as started.
+	c := &Controller{
+		conn:        proxy.clientConn,
+		started:     1,
+		controlAddr: proxy.serverAddr,
+	}
+
+	// Accept the connection inside a goroutine. We will also write some
+	// data so that the reconnection can succeed. We will mock three writes
+	// and two reads inside our proxy server,
+	//   - write protocol info
+	//   - read auth info
+	//   - write auth challenge
+	//   - read auth challenge
+	//   - write OK
+	go func() {
+		// Accept the new connection.
+		server, err := proxy.server.Accept()
+		require.NoError(t, err, "failed to accept")
+
+		// Write the protocol info.
+		resp := "250-PROTOCOLINFO 1\n" +
+			"250-AUTH METHODS=NULL\n" +
+			"250 OK\n"
+		_, err = server.Write([]byte(resp))
+		require.NoErrorf(t, err, "failed to write protocol info")
+
+		// Read the auth info from the client.
+		buf := make([]byte, 65535)
+		_, err = server.Read(buf)
+		require.NoError(t, err)
+
+		// Write the auth challenge.
+		resp = "250 AUTHCHALLENGE SERVERHASH=fake\n"
+		_, err = server.Write([]byte(resp))
+		require.NoErrorf(t, err, "failed to write auth challenge")
+
+		// Read the auth challenge resp from the client.
+		_, err = server.Read(buf)
+		require.NoError(t, err)
+
+		// Write OK resp.
+		resp = "250 OK\n"
+		_, err = server.Write([]byte(resp))
+		require.NoErrorf(t, err, "failed to write response auth")
+	}()
+
+	// Reconnect should succeed.
+	require.NoError(t, c.Reconnect())
+
+	// Check that the old connection is closed.
+	_, err := proxy.clientConn.ReadLine()
+	require.Contains(t, err.Error(), "use of closed network connection")
+
+	// Check that the connection has been updated.
+	require.NotEqual(t, proxy.clientConn, c.conn)
+}