From 9126db177b2ef5b9eb224ed795ce4b4bc2527073 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Fri, 16 Sep 2016 12:42:41 +1000
Subject: [PATCH] Add Config.{Force,PreferNo}Encryption

---
 client.go | 17 +++++++++++------
 config.go |  7 +++++--
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/client.go b/client.go
index 79409d72..c4bf2812 100644
--- a/client.go
+++ b/client.go
@@ -582,7 +582,8 @@ func (cl *Client) establishOutgoingConn(t *Torrent, addr string) (c *connection,
 	if nc == nil {
 		return
 	}
-	c, err = cl.handshakesConnection(nc, t, !cl.config.DisableEncryption, utp)
+	encryptFirst := !cl.config.DisableEncryption && !cl.config.PreferNoEncryption
+	c, err = cl.handshakesConnection(nc, t, encryptFirst, utp)
 	if err != nil {
 		nc.Close()
 		return
@@ -590,12 +591,12 @@ func (cl *Client) establishOutgoingConn(t *Torrent, addr string) (c *connection,
 		return
 	}
 	nc.Close()
-	if cl.config.DisableEncryption {
-		// We already tried without encryption.
+	if cl.config.DisableEncryption || cl.config.ForceEncryption {
+		// There's no alternate encryption case to try.
 		return
 	}
-	// Try again without encryption, using whichever protocol type worked last
-	// time.
+	// Try again with encryption if we didn't earlier, or without if we did,
+	// using whichever protocol type worked last time.
 	if utp {
 		nc, err = cl.dialUTP(addr, t)
 	} else {
@@ -605,7 +606,7 @@ func (cl *Client) establishOutgoingConn(t *Torrent, addr string) (c *connection,
 		err = fmt.Errorf("error dialing for unencrypted connection: %s", err)
 		return
 	}
-	c, err = cl.handshakesConnection(nc, t, false, utp)
+	c, err = cl.handshakesConnection(nc, t, !encryptFirst, utp)
 	if err != nil || c == nil {
 		nc.Close()
 	}
@@ -852,6 +853,10 @@ func (cl *Client) receiveHandshakes(c *connection) (t *Torrent, err error) {
 			return
 		}
 	}
+	if cl.config.ForceEncryption && !c.encrypted {
+		err = errors.New("connection not encrypted")
+		return
+	}
 	ih, ok, err := cl.connBTHandshake(c, nil)
 	if err != nil {
 		err = fmt.Errorf("error during bt handshake: %s", err)
diff --git a/config.go b/config.go
index 8b180fdc..dbcdbd13 100644
--- a/config.go
+++ b/config.go
@@ -36,8 +36,11 @@ type Config struct {
 	// Called to instantiate storage for each added torrent. Builtin backends
 	// are in the storage package. If not set, the "file" implementation is
 	// used.
-	DefaultStorage    storage.ClientImpl
-	DisableEncryption bool `long:"disable-encryption"`
+	DefaultStorage storage.ClientImpl
+
+	DisableEncryption  bool `long:"disable-encryption"`
+	ForceEncryption    bool // Don't allow unobfuscated connections.
+	PreferNoEncryption bool
 
 	IPBlocklist iplist.Ranger
 	DisableIPv6 bool `long:"disable-ipv6"`