瀏覽代碼

Add ability to write and parse push promise frames.

A lot of the functionality in push promises are similar to headers
frames, so some of the code has been copied and adapted. Waiting
until there's a user of the push promise frames to dedup the code.
Daniel Morsing 11 年之前
父節點
當前提交
b7c97da766
共有 2 個文件被更改,包括 152 次插入1 次删除
  1. 114 1
      frame.go
  2. 38 0
      frame_test.go

+ 114 - 1
frame.go

@@ -84,6 +84,9 @@ const (
 
 	// Continuation Frame
 	FlagContinuationEndHeaders Flags = 0x4
+
+	FlagPushPromiseEndHeaders = 0x4
+	FlagPushPromisePadded     = 0x8
 )
 
 var flagName = map[FrameType]map[Flags]string{
@@ -106,6 +109,10 @@ var flagName = map[FrameType]map[Flags]string{
 	FrameContinuation: {
 		FlagContinuationEndHeaders: "END_HEADERS",
 	},
+	FramePushPromise: {
+		FlagPushPromiseEndHeaders: "END_HEADERS",
+		FlagPushPromisePadded:     "PADDED",
+	},
 }
 
 // a frameParser parses a frame given its FrameHeader and payload
@@ -119,7 +126,7 @@ var frameParsers = map[FrameType]frameParser{
 	FramePriority:     parsePriorityFrame,
 	FrameRSTStream:    parseRSTStreamFrame,
 	FrameSettings:     parseSettingsFrame,
-	FramePushPromise:  nil, // TODO
+	FramePushPromise:  parsePushPromise,
 	FramePing:         parsePingFrame,
 	FrameGoAway:       parseGoAwayFrame,
 	FrameWindowUpdate: parseWindowUpdateFrame,
@@ -951,6 +958,112 @@ func (f *Framer) WriteContinuation(streamID uint32, endHeaders bool, headerBlock
 	return f.endWrite()
 }
 
+// A PushPromiseFrame is used to initiate a server stream.
+// See http://http2.github.io/http2-spec/#rfc.section.6.6
+type PushPromiseFrame struct {
+	FrameHeader
+	PromiseID     uint32
+	headerFragBuf []byte // not owned
+}
+
+func (f *PushPromiseFrame) HeaderBlockFragment() []byte {
+	f.checkValid()
+	return f.headerFragBuf
+}
+
+func (f *PushPromiseFrame) HeadersEnded() bool {
+	return f.FrameHeader.Flags.Has(FlagPushPromiseEndHeaders)
+}
+
+func parsePushPromise(fh FrameHeader, p []byte) (_ Frame, err error) {
+	pp := &PushPromiseFrame{
+		FrameHeader: fh,
+	}
+	if pp.StreamID == 0 {
+		// PUSH_PROMISE frames MUST be associated with an existing,
+		// peer-initiated stream. The stream identifier of a
+		// PUSH_PROMISE frame indicates the stream it is associated
+		// with. If the stream identifier field specifies the value
+		// 0x0, a recipient MUST respond with a connection error
+		// (Section 5.4.1) of type PROTOCOL_ERROR.
+		return nil, ConnectionError(ErrCodeProtocol)
+	}
+	// The PUSH_PROMISE frame includes optional padding.
+	// Padding fields and flags are identical to those defined for DATA frames
+	var padLength uint8
+	if fh.Flags.Has(FlagPushPromisePadded) {
+		if p, padLength, err = readByte(p); err != nil {
+			return
+		}
+	}
+
+	p, pp.PromiseID, err = readUint32(p)
+	if err != nil {
+		return
+	}
+	pp.PromiseID = pp.PromiseID & (1<<31 - 1)
+
+	if int(padLength) > len(p) {
+		// like the DATA frame, error out if padding is longer than the body.
+		return nil, ConnectionError(ErrCodeProtocol)
+	}
+	pp.headerFragBuf = p[:len(p)-int(padLength)]
+	return pp, nil
+}
+
+// PushPromiseParam are the parameters for writing a PUSH_PROMISE frame.
+type PushPromiseParam struct {
+	// StreamID is the required Stream ID to initiate.
+	StreamID uint32
+
+	// PromiseID is the required Stream ID which this
+	// Push Promises
+	PromiseID uint32
+
+	// BlockFragment is part (or all) of a Header Block.
+	BlockFragment []byte
+
+	// EndHeaders indicates that this frame contains an entire
+	// header block and is not followed by any
+	// CONTINUATION frames.
+	EndHeaders bool
+
+	// PadLength is the optional number of bytes of zeros to add
+	// to this frame.
+	PadLength uint8
+}
+
+// WritePushPromise writes a single PushPromise Frame.
+//
+// As with Header Frames, This is the low level call for writing
+// individual frames. Continuation frames are handled elsewhere.
+//
+// It will perform exactly one Write to the underlying Writer.
+// It is the caller's responsibility to not call other Write methods concurrently.
+func (f *Framer) WritePushPromise(p PushPromiseParam) error {
+	if !validStreamID(p.StreamID) && !f.AllowIllegalWrites {
+		return errStreamID
+	}
+	var flags Flags
+	if p.PadLength != 0 {
+		flags |= FlagPushPromisePadded
+	}
+	if p.EndHeaders {
+		flags |= FlagPushPromiseEndHeaders
+	}
+	f.startWrite(FramePushPromise, flags, p.StreamID)
+	if p.PadLength != 0 {
+		f.writeByte(p.PadLength)
+	}
+	if !validStreamID(p.PromiseID) && !f.AllowIllegalWrites {
+		return errStreamID
+	}
+	f.writeUint32(p.PromiseID)
+	f.wbuf = append(f.wbuf, p.BlockFragment...)
+	f.wbuf = append(f.wbuf, padZeros[:p.PadLength]...)
+	return f.endWrite()
+}
+
 // WriteRawFrame writes a raw frame. This can be used to write
 // extension frames unknown to this package.
 func (f *Framer) WriteRawFrame(t FrameType, flags Flags, streamID uint32, payload []byte) error {

+ 38 - 0
frame_test.go

@@ -501,3 +501,41 @@ func TestWriteGoAway(t *testing.T) {
 		t.Errorf("debug data = %q; want %q", got, debug)
 	}
 }
+
+func TestWritePushPromise(t *testing.T) {
+	pp := PushPromiseParam{
+		StreamID:      42,
+		PromiseID:     42,
+		BlockFragment: []byte("abc"),
+	}
+	fr, buf := testFramer()
+	if err := fr.WritePushPromise(pp); err != nil {
+		t.Fatal(err)
+	}
+	const wantEnc = "\x00\x00\x07\x05\x00\x00\x00\x00*\x00\x00\x00*abc"
+	if buf.String() != wantEnc {
+		t.Errorf("encoded as %q; want %q", buf.Bytes(), wantEnc)
+	}
+	f, err := fr.ReadFrame()
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, ok := f.(*PushPromiseFrame)
+	if !ok {
+		t.Fatalf("got %T; want *PushPromiseFrame", f)
+	}
+	want := &PushPromiseFrame{
+		FrameHeader: FrameHeader{
+			valid:    true,
+			Type:     0x5,
+			Flags:    0x0,
+			Length:   0x7,
+			StreamID: 42,
+		},
+		PromiseID:     42,
+		headerFragBuf: []byte("abc"),
+	}
+	if !reflect.DeepEqual(f, want) {
+		t.Fatalf("parsed back:\n%#v\nwant:\n%#v", f, want)
+	}
+}