From c9560c2f9af69d340417f1bf647a2777b6283d7d Mon Sep 17 00:00:00 2001
From: Kaleb Elwert <belak@coded.io>
Date: Tue, 12 Dec 2017 13:14:48 -0800
Subject: [PATCH] Add documentation and error if a CAP is requested outside the
 handshake

---
 client.go      | 17 ++++++++++++++++-
 client_test.go |  7 +++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/client.go b/client.go
index b10d78e..d686823 100644
--- a/client.go
+++ b/client.go
@@ -280,17 +280,33 @@ func (c *Client) sendError(err error) {
 	}
 }
 
+// CapRequest allows you to request IRCv3 capabilities from the server during
+// the handshake. This must be called before the handshake completes and so it
+// is recommended that this be called before Run. If the CAP is marked as
+// required, the client will exit if that CAP could not be negotiated during the
+// handshake.
 func (c *Client) CapRequest(capName string, required bool) {
+	if c.finishedHandshake {
+		c.sendError(errors.New("CAP requested after CAP handshake"))
+		return
+	}
+
 	cap := c.caps[capName]
 	cap.Requested = true
 	cap.Required = cap.Required || required
 	c.caps[capName] = cap
 }
 
+// CapEnabled allows you to check if a CAP is enabled for this connection. Note
+// that it will not be populated until after the CAP handshake is done, so it is
+// recommended to wait to check this until after a message like 001.
 func (c *Client) CapEnabled(capName string) bool {
 	return c.caps[capName].Enabled
 }
 
+// CapAvailable allows you to check if a CAP is available on this server. Note
+// that it will not be populated until after the CAP handshake is done, so it is
+// recommended to wait to check this until after a message like 001.
 func (c *Client) CapAvailable(capName string) bool {
 	return c.caps[capName].Available
 }
@@ -341,7 +357,6 @@ func (c *Client) Run() error {
 
 		m, err := c.ReadMessage()
 		if err != nil {
-			fmt.Println("READ EOF")
 			c.sendError(err)
 			break
 		}
diff --git a/client_test.go b/client_test.go
index d608a90..77e2aa3 100644
--- a/client_test.go
+++ b/client_test.go
@@ -149,6 +149,13 @@ func TestCapReq(t *testing.T) {
 	assert.False(t, c.CapEnabled("multi-prefix"))
 	assert.False(t, c.CapAvailable("random-thing"))
 	assert.True(t, c.CapAvailable("multi-prefix"))
+
+	runClientTest(t, config, errors.New("CAP requested after CAP handshake"), func(c *Client) {
+		assert.False(t, c.CapAvailable("random-thing"))
+		assert.False(t, c.CapAvailable("multi-prefix"))
+		c.finishedHandshake = true
+		c.CapRequest("multi-prefix", true)
+	}, []TestAction{})
 }
 
 func TestClient(t *testing.T) {