Merge pull request #915 from slingamn/issue817_fakelag.1

fix #817

irc/client.go

@@ -112,9 +112,10 @@ type Session struct {
 	rawHostname string
 	isTor       bool
 
-	idletimer IdleTimer
-	fakelag   Fakelag
-	destroyed uint32
+	idletimer            IdleTimer
+	fakelag              Fakelag
+	deferredFakelagCount int
+	destroyed            uint32
 
 	certfp          string
 	sasl            saslStatus
@@ -148,6 +149,42 @@ type MultilineBatch struct {
 	tags          map[string]string
 }
 
+// Starts a multiline batch, failing if there's one already open
+func (s *Session) StartMultilineBatch(label, target, responseLabel string, tags map[string]string) (err error) {
+	if s.batch.label != "" {
+		return errInvalidMultilineBatch
+	}
+
+	s.batch.label, s.batch.target, s.batch.responseLabel, s.batch.tags = label, target, responseLabel, tags
+	s.fakelag.Suspend()
+	return
+}
+
+// Closes a multiline batch unconditionally; returns the batch and whether
+// it was validly terminated (pass "" as the label if you don't care about the batch)
+func (s *Session) EndMultilineBatch(label string) (batch MultilineBatch, err error) {
+	batch = s.batch
+	s.batch = MultilineBatch{}
+	s.fakelag.Unsuspend()
+
+	// heuristics to estimate how much data they used while fakelag was suspended
+	fakelagBill := (batch.message.LenBytes() / 512) + 1
+	fakelagBillLines := (batch.message.LenLines() * 60) / 512
+	if fakelagBill < fakelagBillLines {
+		fakelagBill = fakelagBillLines
+	}
+	s.deferredFakelagCount = fakelagBill
+
+	if batch.label == "" || batch.label != label || batch.message.LenLines() == 0 {
+		err = errInvalidMultilineBatch
+		return
+	}
+
+	batch.message.SetTime()
+
+	return
+}
+
 // sets the session quit message, if there isn't one already
 func (sd *Session) SetQuitMessage(message string) (set bool) {
 	if message == "" {
@@ -596,7 +633,11 @@ func (client *Client) run(session *Session, proxyLine string) {
 		}
 
 		if client.registered {
-			session.fakelag.Touch()
+			touches := session.deferredFakelagCount + 1
+			session.deferredFakelagCount = 0
+			for i := 0; i < touches; i++ {
+				session.fakelag.Touch()
+			}
 		} else {
 			// DoS hardening, #505
 			session.registrationMessages++
@@ -617,19 +658,6 @@ func (client *Client) run(session *Session, proxyLine string) {
 			break
 		}
 
-		// "Clients MUST NOT send messages other than PRIVMSG while a multiline batch is open."
-		// in future we might want to whitelist some commands that are allowed here, like PONG
-		if session.batch.label != "" && msg.Command != "BATCH" {
-			_, batchTag := msg.GetTag("batch")
-			if batchTag != session.batch.label {
-				if msg.Command != "NOTICE" {
-					session.Send(nil, client.server.name, "FAIL", "BATCH", "MULTILINE_INVALID", client.t("Incorrect batch tag sent"))
-				}
-				session.batch = MultilineBatch{}
-				continue
-			}
-		}
-
 		cmd, exists := Commands[msg.Command]
 		if !exists {
 			if len(msg.Command) > 0 {

irc/commands.go

@@ -47,7 +47,7 @@ func (cmd *Command) Run(server *Server, client *Client, session *Session, msg ir
 		}
 		if session.batch.label != "" && !cmd.allowedInBatch {
 			rb.Add(nil, server.name, "FAIL", "BATCH", "MULTILINE_INVALID", client.t("Command not allowed during a multiline batch"))
-			session.batch = MultilineBatch{}
+			session.EndMultilineBatch("")
 			return false
 		}
 

irc/errors.go

@@ -60,6 +60,7 @@ var (
 	errCASFailed                      = errors.New("Compare-and-swap update of database value failed")
 	errEmptyCredentials               = errors.New("No more credentials are approved")
 	errCredsExternallyManaged         = errors.New("Credentials are externally managed and cannot be changed here")
+	errInvalidMultilineBatch          = errors.New("Invalid multiline batch")
 )
 
 // Socket Errors

irc/fakelag.go

@@ -25,6 +25,7 @@ const (
 // from the loop that accepts the client's input and runs commands
 type Fakelag struct {
 	config    FakelagConfig
+	suspended bool
 	nowFunc   func() time.Time
 	sleepFunc func(time.Duration)
 
@@ -40,6 +41,22 @@ func (fl *Fakelag) Initialize(config FakelagConfig) {
 	fl.state = FakelagBursting
 }
 
+// Idempotently turn off fakelag if it's enabled
+func (fl *Fakelag) Suspend() {
+	if fl.config.Enabled {
+		fl.suspended = true
+		fl.config.Enabled = false
+	}
+}
+
+// Idempotently turn fakelag back on if it was previously Suspend'ed
+func (fl *Fakelag) Unsuspend() {
+	if fl.suspended {
+		fl.config.Enabled = true
+		fl.suspended = false
+	}
+}
+
 // register a new command, sleep if necessary to delay it
 func (fl *Fakelag) Touch() {
 	if !fl.config.Enabled {

irc/fakelag_test.go

@@ -121,3 +121,35 @@ func TestFakelag(t *testing.T) {
 		t.Fatalf("should not have slept")
 	}
 }
+
+func TestSuspend(t *testing.T) {
+	window, _ := time.ParseDuration("1s")
+	fl, _ := newFakelagForTesting(window, 3, 2, window)
+	assertEqual(fl.config.Enabled, true, t)
+
+	// suspend idempotently disables
+	fl.Suspend()
+	assertEqual(fl.config.Enabled, false, t)
+	fl.Suspend()
+	assertEqual(fl.config.Enabled, false, t)
+	// unsuspend idempotently enables
+	fl.Unsuspend()
+	assertEqual(fl.config.Enabled, true, t)
+	fl.Unsuspend()
+	assertEqual(fl.config.Enabled, true, t)
+	fl.Suspend()
+	assertEqual(fl.config.Enabled, false, t)
+
+	fl2, _ := newFakelagForTesting(window, 3, 2, window)
+	fl2.config.Enabled = false
+
+	// if we were never enabled, suspend and unsuspend are both no-ops
+	fl2.Suspend()
+	assertEqual(fl2.config.Enabled, false, t)
+	fl2.Suspend()
+	assertEqual(fl2.config.Enabled, false, t)
+	fl2.Unsuspend()
+	assertEqual(fl2.config.Enabled, false, t)
+	fl2.Unsuspend()
+	assertEqual(fl2.config.Enabled, false, t)
+}

irc/handlers.go

@@ -332,30 +332,20 @@ func batchHandler(server *Server, client *Client, msg ircmsg.IrcMessage, rb *Res
 	if len(tag) == 0 {
 		fail = true
 	} else if tag[0] == '+' {
-		if rb.session.batch.label != "" || msg.Params[1] != caps.MultilineBatchType {
+		if len(msg.Params) < 3 || msg.Params[1] != caps.MultilineBatchType {
 			fail = true
 		} else {
-			rb.session.batch.label = tag[1:]
-			rb.session.batch.tags = msg.ClientOnlyTags()
-			if len(msg.Params) == 2 {
-				fail = true
-			} else {
-				rb.session.batch.target = msg.Params[2]
-				// save the response label for later
-				rb.session.batch.responseLabel = rb.Label
+			err := rb.session.StartMultilineBatch(tag[1:], msg.Params[2], rb.Label, msg.ClientOnlyTags())
+			fail = (err != nil)
+			if !fail {
+				// suppress ACK for the initial BATCH message (we'll apply the stored label later)
 				rb.Label = ""
 			}
 		}
 	} else if tag[0] == '-' {
-		if rb.session.batch.label == "" || rb.session.batch.label != tag[1:] {
-			fail = true
-		} else if rb.session.batch.message.LenLines() == 0 {
-			fail = true
-		} else {
-			batch := rb.session.batch
-			rb.session.batch = MultilineBatch{}
-			// time tag should correspond to the time when the message was completed
-			batch.message.SetTime()
+		batch, err := rb.session.EndMultilineBatch(tag[1:])
+		fail = (err != nil)
+		if !fail {
 			histType, err := msgCommandToHistType(batch.command)
 			if err != nil {
 				histType = history.Privmsg
@@ -369,7 +359,7 @@ func batchHandler(server *Server, client *Client, msg ircmsg.IrcMessage, rb *Res
 	}
 
 	if fail {
-		rb.session.batch = MultilineBatch{}
+		rb.session.EndMultilineBatch("")
 		if sendErrors {
 			rb.Add(nil, server.name, "FAIL", "BATCH", "MULTILINE_INVALID", client.t("Invalid multiline batch"))
 		}
@@ -1813,9 +1803,17 @@ func nickHandler(server *Server, client *Client, msg ircmsg.IrcMessage, rb *Resp
 
 // helper to store a batched PRIVMSG in the session object
 func absorbBatchedMessage(server *Server, client *Client, msg ircmsg.IrcMessage, batchTag string, histType history.ItemType, rb *ResponseBuffer) {
-	// sanity checks. batch tag correctness was already checked and is redundant here
-	// as a defensive measure. TAGMSG is checked without an error message: "don't eat paste"
-	if batchTag != rb.session.batch.label || histType == history.Tagmsg || len(msg.Params) == 1 || msg.Params[1] == "" {
+	if batchTag != rb.session.batch.label {
+		if histType != history.Notice {
+			rb.Add(nil, server.name, "FAIL", "BATCH", "MULTILINE_INVALID", client.t("Incorrect batch tag sent"))
+		}
+		rb.session.EndMultilineBatch("")
+		return
+	} else if len(msg.Params) < 2 || msg.Params[1] == "" {
+		if histType != history.Notice {
+			rb.Add(nil, server.name, "FAIL", "BATCH", "MULTILINE_INVALID", client.t("Invalid multiline batch"))
+		}
+		rb.session.EndMultilineBatch("")
 		return
 	}
 	rb.session.batch.command = msg.Command
@@ -1826,12 +1824,12 @@ func absorbBatchedMessage(server *Server, client *Client, msg ircmsg.IrcMessage,
 		if histType != history.Notice {
 			rb.Add(nil, server.name, "FAIL", "BATCH", "MULTILINE_MAX_BYTES", strconv.Itoa(config.Limits.Multiline.MaxBytes))
 		}
-		rb.session.batch = MultilineBatch{}
+		rb.session.EndMultilineBatch("")
 	} else if config.Limits.Multiline.MaxLines != 0 && config.Limits.Multiline.MaxLines < rb.session.batch.message.LenLines() {
 		if histType != history.Notice {
 			rb.Add(nil, server.name, "FAIL", "BATCH", "MULTILINE_MAX_LINES", strconv.Itoa(config.Limits.Multiline.MaxLines))
 		}
-		rb.session.batch = MultilineBatch{}
+		rb.session.EndMultilineBatch("")
 	}
 }