package network import ( "context" "crypto/tls" "errors" "fmt" "git.hpds.cc/Component/network/auth" "git.hpds.cc/Component/network/frame" "github.com/quic-go/quic-go" ) // ControlStream defines the interface for controlling a stream. type ControlStream interface { // CloseStream notifies the peer's control stream to close the data stream with the given streamID and error message. CloseStream(streamId string, errString string) error // ReceiveStreamClose is received from the peer's control stream to close the data stream according to streamID and error message. ReceiveStreamClose() (streamId string, errString string, err error) // CloseWithError closes the control stream. CloseWithError(code uint64, errString string) error } // ServerControlStream defines the interface of server side control stream. type ServerControlStream interface { ControlStream // VerifyAuthentication verify the Authentication from client side. VerifyAuthentication(verifyFunc func(auth.Object) (bool, error)) error // AcceptStream accepts data stream from the request of client. AcceptStream(context.Context) (DataStream, error) } // ClientControlStream defines the interface of client side control stream. type ClientControlStream interface { ControlStream // Authenticate with credential, the credential will be sent to ServerControlStream to authenticate the client. Authenticate(*auth.Credential) error // OpenStream request a ServerControlStream to create a new data stream. OpenStream(context.Context, *frame.HandshakeFrame) (DataStream, error) } var _ ServerControlStream = &serverControlStream{} type serverControlStream struct { conn quic.Connection stream frame.ReadWriter } // NewServerControlStream returns ServerControlStream from quic Connection and the first stream of this Connection. func NewServerControlStream(qConn quic.Connection, stream frame.ReadWriter) ServerControlStream { return &serverControlStream{ conn: qConn, stream: stream, } } func (ss *serverControlStream) ReceiveStreamClose() (streamId string, errReason string, err error) { return receiveStreamClose(ss.stream) } func (ss *serverControlStream) CloseStream(streamId string, errString string) error { return closeStream(ss.stream, streamId, errString) } func (ss *serverControlStream) AcceptStream(context.Context) (DataStream, error) { f, err := ss.stream.ReadFrame() if err != nil { return nil, err } switch ff := f.(type) { case *frame.HandshakeFrame: stream, err := ss.conn.OpenStreamSync(context.Background()) if err != nil { return nil, err } _, err = stream.Write(frame.NewHandshakeAckFrame(ff.ID()).Encode()) if err != nil { return nil, err } dataStream := newDataStream( ff.Name(), ff.ID(), StreamType(ff.StreamType()), ff.Metadata(), stream, ff.ObserveDataTags(), ss, ) return dataStream, nil default: return nil, fmt.Errorf("yomo: control stream read unexpected frame %s", f.Type()) } } func (ss *serverControlStream) CloseWithError(code uint64, errString string) error { return closeWithError(ss.conn, code, errString) } func (ss *serverControlStream) VerifyAuthentication(verifyFunc func(auth.Object) (bool, error)) error { first, err := ss.stream.ReadFrame() if err != nil { return err } received, ok := first.(*frame.AuthenticationFrame) if !ok { return fmt.Errorf("yomo: read unexpected frame while waiting for authentication, frame read: %s", received.Type().String()) } ok, err = verifyFunc(received) if err != nil { return err } if !ok { return ss.stream.WriteFrame( frame.NewAuthenticationRespFrame( false, fmt.Sprintf("yomo: authentication failed, client credential name is %s", received.AuthName()), ), ) } return ss.stream.WriteFrame(frame.NewAuthenticationRespFrame(true, "")) } var _ ClientControlStream = &clientControlStream{} type clientControlStream struct { conn quic.Connection stream frame.ReadWriter } // OpenClientControlStream opens ClientControlStream from addr. func OpenClientControlStream( ctx context.Context, addr string, tlsConfig *tls.Config, quicConfig *quic.Config, ) (ClientControlStream, error) { conn, err := quic.DialAddrContext(ctx, addr, tlsConfig, quicConfig) if err != nil { return nil, err } stream, err := conn.OpenStream() if err != nil { return nil, err } return NewClientControlStream(conn, NewFrameStream(stream)), nil } // NewClientControlStream returns ClientControlStream from quic Connection and the first stream form the Connection. func NewClientControlStream(qConn quic.Connection, stream frame.ReadWriter) ClientControlStream { return &clientControlStream{ conn: qConn, stream: stream, } } func (cs *clientControlStream) ReceiveStreamClose() (streamId string, errReason string, err error) { return receiveStreamClose(cs.stream) } func (cs *clientControlStream) CloseStream(streamId string, errString string) error { return closeStream(cs.stream, streamId, errString) } func (cs *clientControlStream) Authenticate(cred *auth.Credential) error { if err := cs.stream.WriteFrame( frame.NewAuthenticationFrame(cred.Name(), cred.Payload())); err != nil { return err } received, err := cs.stream.ReadFrame() if err != nil { return err } resp, ok := received.(*frame.AuthenticationRespFrame) if !ok { return fmt.Errorf( "yomo: read unexcept frame during waiting authentication resp, frame readed: %s", received.Type().String(), ) } if !resp.OK() { return errors.New(resp.Reason()) } return nil } // dataStreamAcked drain HandshakeAckFrame from stream. func dataStreamAcked(stream DataStream) error { first, err := stream.ReadFrame() if err != nil { return err } f, ok := first.(*frame.HandshakeAckFrame) if !ok { return fmt.Errorf("yomo: data stream read first frame should be HandshakeAckFrame, but got %s", first.Type().String()) } if f.StreamId() != stream.ID() { return fmt.Errorf("yomo: data stream ack exception, stream id did not match") } return nil } func (cs *clientControlStream) OpenStream(ctx context.Context, hf *frame.HandshakeFrame) (DataStream, error) { err := cs.stream.WriteFrame(frame.NewHandshakeFrame( hf.Name(), hf.ID(), hf.StreamType(), hf.ObserveDataTags(), hf.Metadata(), )) if err != nil { return nil, err } quicStream, err := cs.conn.AcceptStream(ctx) if err != nil { return nil, err } dataStream := newDataStream( hf.Name(), hf.ID(), StreamType(hf.StreamType()), hf.Metadata(), quicStream, hf.ObserveDataTags(), cs, ) if err := dataStreamAcked(dataStream); err != nil { return nil, err } return dataStream, nil } func (cs *clientControlStream) CloseWithError(code uint64, errString string) error { return closeWithError(cs.conn, code, errString) } func closeStream(controlStream frame.Writer, streamID string, errString string) error { f := frame.NewCloseStreamFrame(streamID, errString) return controlStream.WriteFrame(f) } func receiveStreamClose(controlStream frame.Reader) (streamID string, errString string, err error) { f, err := controlStream.ReadFrame() if err != nil { return "", "", err } ff, ok := f.(*frame.CloseStreamFrame) if !ok { return "", "", errors.New("yomo: control stream only transmit close stream frame") } return ff.StreamID(), ff.Reason(), nil } func closeWithError(qConn quic.Connection, code uint64, errString string) error { return qConn.CloseWithError( quic.ApplicationErrorCode(code), errString, ) }