Fix a bug in writeTrailer where EOF never got written to the response

twitchtv · Jun 21, 2018 · b8fadd5 · b8fadd5
1 parent fe19e01
commit b8fadd5
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/protoc-gen-twirp/generator.go b/protoc-gen-twirp/generator.go
@@ -1440,6 +1440,13 @@ func (t *twirp) generateServerProtobufMethod(service *descriptor.ServiceDescript
 		t.P(`	trailer := `, t.pkgs["proto"], `.NewBuffer(nil)`)
 		t.P(`	_ = trailer.EncodeVarint((2 << 3) | 2) // field tag`)
 		t.P(`	writeTrailer := func(err error) {`)
+		t.P(`		defer func() {`)
+		t.P(`			_, writeErr := resp.Write(trailer.Bytes())`)
+		t.P(`			if writeErr != nil {`)
+		t.P(`				// Ignored, for the same reason as in the writeError func`)
+		t.P(`				_ = writeErr`)
+		t.P(`			}`)
+		t.P(`		}()`)
 		t.P(`		if err == `, t.pkgs["io"], `.EOF {`)
 		t.P(`			trailer.EncodeStringBytes("EOF")`)
 		t.P(`			return`)
@@ -1455,11 +1462,6 @@ func (t *twirp) generateServerProtobufMethod(service *descriptor.ServiceDescript
 		t.P(`		if encodeErr := trailer.EncodeStringBytes(string(marshalErrorToJSON(twerr))); encodeErr != nil {`)
 		t.P(`			_ = trailer.EncodeStringBytes("{\"code\":\"" + string(`, t.pkgs["twirp"], `.Internal) + "\",\"msg\":\"There was an error but it could not be serialized into JSON\"}") // fallback`)
 		t.P(`		}`)
-		t.P(`		_, writeErr := resp.Write(trailer.Bytes())`)
-		t.P(`		if writeErr != nil {`)
-		t.P(`			// Ignored, for the same reason as in the writeError func`)
-		t.P(`			_ = writeErr`)
-		t.P(`		}`)
 		t.P(`	}`)
 		t.P(``)
 		t.P(`	messages := `, t.pkgs["proto"], `.NewBuffer(nil)`)