mirror of
https://github.com/soheilhy/cmux.git
synced 2025-01-18 18:56:26 +08:00
Initial commit
This commit is contained in:
Soheil Hassas Yeganeh
commit
754f5b897d
24
.gitignore
vendored
Normal file
24
.gitignore
vendored
Normal file
@ -0,0 +1,24 @@
|
||||
# Compiled Object files, Static and Dynamic libs (Shared Objects)
|
||||
*.o
|
||||
*.a
|
||||
*.so
|
||||
|
||||
# Folders
|
||||
_obj
|
||||
_test
|
||||
|
||||
# Architecture specific extensions/prefixes
|
||||
*.[568vq]
|
||||
[568vq].out
|
||||
|
||||
*.cgo1.go
|
||||
*.cgo2.c
|
||||
_cgo_defun.c
|
||||
_cgo_gotypes.go
|
||||
_cgo_export.*
|
||||
|
||||
_testmain.go
|
||||
|
||||
*.exe
|
||||
*.test
|
||||
*.prof
|
||||
5
.travis.yml
Normal file
5
.travis.yml
Normal file
@ -0,0 +1,5 @@
|
||||
language: go
|
||||
|
||||
go:
|
||||
- 1.3
|
||||
- 1.4
|
||||
57
README.md
Normal file
57
README.md
Normal file
@ -0,0 +1,57 @@
|
||||
# cmux: Connection Mux  
|
||||
cmux is a generic Go library to multiplex connections based on
|
||||
their content. Using cmux, one can serve gRPC, HTTP, and Go RPC
|
||||
on the same TCP listener to avoid having to use one port per
|
||||
protocol.
|
||||
|
||||
## How-To
|
||||
Simply create your main listener, create a cmux for that listener,
|
||||
and then match connections:
|
||||
```go
|
||||
// Create the main listener.
|
||||
l, err := net.Listen("tcp", ":23456")
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
// Create a cmux.
|
||||
m := cmux.New(l)
|
||||
|
||||
// Match connections in order.
|
||||
grpcl := m.Match(cmux.HTTP2HeaderField("content-type", "application/grpc"))
|
||||
httpl := m.Match(cmux.Any()) // Any means anything that is not yet matched.
|
||||
|
||||
// Create your protocol servers.
|
||||
grpcS := grpc.NewServer()
|
||||
pb.RegisterGreeterServer(grpcs, &server{})
|
||||
|
||||
httpS := &http.Server{
|
||||
Handler: &testHTTP1Handler{},
|
||||
}
|
||||
|
||||
// Use the muxed listeners for your servers.
|
||||
go grpcS.Serve(grpcl)
|
||||
go httpS.Serve(httpl)
|
||||
|
||||
// Start serving!
|
||||
m.Serve()
|
||||
```
|
||||
|
||||
Take a look at [other examples in the GoDoc](http://localhost:6060/pkg/github.com/soheilhy/cmux/#pkg-examples).
|
||||
|
||||
## Docs
|
||||
* [GoDocs](https://godoc.org/github.com/soheilhy/cmux)
|
||||
|
||||
## Performance
|
||||
There is a huge room for improvment but since we are only matching
|
||||
the very first bytes of a connection, the performance overheads on
|
||||
long-lived connections (i.e., RPCs and pipelined HTTP streams)
|
||||
is negligible.
|
||||
|
||||
*TODO(soheil)*: Add benchmarks.
|
||||
|
||||
## Limitations
|
||||
*TLS*: Since `cmux` sits in between the actual listener and the mux'ed
|
||||
listeners, TLS handshake is not handled inside the actual servers.
|
||||
Because of that, when you handle HTTPS using cmux `http.Request.TLS`
|
||||
would not be set.
|
||||
197
cmux.go
Normal file
197
cmux.go
Normal file
@ -0,0 +1,197 @@
|
||||
package cmux
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
)
|
||||
|
||||
// Matcher matches a connection based on its content.
|
||||
type Matcher func(r io.Reader) (ok bool)
|
||||
|
||||
// ErrorHandler handles an error and returns whether
|
||||
// the mux should continue serving the listener.
|
||||
type ErrorHandler func(err error) (ok bool)
|
||||
|
||||
// ErrNotMatched is returned whenever a connection is not matched by any of
|
||||
// the matchers registered in the multiplexer.
|
||||
type ErrNotMatched struct {
|
||||
c net.Conn
|
||||
}
|
||||
|
||||
func (e ErrNotMatched) Error() string {
|
||||
return fmt.Sprintf("mux: connection %v not matched by an matcher",
|
||||
e.c.RemoteAddr())
|
||||
}
|
||||
|
||||
func (e ErrNotMatched) Temporary() bool { return true }
|
||||
func (e ErrNotMatched) Timeout() bool { return false }
|
||||
|
||||
type errListenerClosed string
|
||||
|
||||
func (e errListenerClosed) Error() string { return string(e) }
|
||||
func (e errListenerClosed) Temporary() bool { return false }
|
||||
func (e errListenerClosed) Timeout() bool { return false }
|
||||
|
||||
var (
|
||||
ErrListenerClosed = errListenerClosed("mux: listener closed")
|
||||
)
|
||||
|
||||
// New instantiates a new connection multiplexer.
|
||||
func New(l net.Listener) CMux {
|
||||
if !flag.Parsed() {
|
||||
flag.Parse()
|
||||
}
|
||||
|
||||
return &cMux{
|
||||
root: l,
|
||||
bufLen: 1024,
|
||||
errh: func(err error) bool { return true },
|
||||
}
|
||||
}
|
||||
|
||||
// CMux is a multiplexer for network connections.
|
||||
type CMux interface {
|
||||
// Match returns a net.Listener that sees (i.e., accepts) only
|
||||
// the connections matched by at least one of the matcher.
|
||||
//
|
||||
// The order used to call Match determines the priority of matchers.
|
||||
Match(matchers ...Matcher) net.Listener
|
||||
// Serve starts multiplexing the listener. Serve blocks and perhaps
|
||||
// should be invoked concurrently within a go routine.
|
||||
Serve() error
|
||||
// HandleError registers an error handler that handles listener errors.
|
||||
HandleError(h ErrorHandler)
|
||||
}
|
||||
|
||||
type matchersListener struct {
|
||||
ss []Matcher
|
||||
l muxListener
|
||||
}
|
||||
|
||||
type cMux struct {
|
||||
root net.Listener
|
||||
bufLen int
|
||||
errh ErrorHandler
|
||||
sls []matchersListener
|
||||
}
|
||||
|
||||
func (m *cMux) Match(matchers ...Matcher) (l net.Listener) {
|
||||
ml := muxListener{
|
||||
Listener: m.root,
|
||||
cch: make(chan net.Conn, m.bufLen),
|
||||
}
|
||||
m.sls = append(m.sls, matchersListener{ss: matchers, l: ml})
|
||||
return ml
|
||||
}
|
||||
|
||||
func (m *cMux) Serve() error {
|
||||
defer func() {
|
||||
for _, sl := range m.sls {
|
||||
close(sl.l.cch)
|
||||
}
|
||||
}()
|
||||
|
||||
for {
|
||||
c, err := m.root.Accept()
|
||||
if err != nil {
|
||||
if !m.handleErr(err) {
|
||||
return err
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
muc := newMuxConn(c)
|
||||
matched := false
|
||||
outer:
|
||||
for _, sl := range m.sls {
|
||||
for _, s := range sl.ss {
|
||||
matched = s(muc.sniffer())
|
||||
muc.reset()
|
||||
if matched {
|
||||
sl.l.cch <- muc
|
||||
break outer
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !matched {
|
||||
c.Close()
|
||||
err := ErrNotMatched{c: c}
|
||||
if !m.handleErr(err) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (m *cMux) HandleError(h ErrorHandler) {
|
||||
m.errh = h
|
||||
}
|
||||
|
||||
func (m *cMux) handleErr(err error) bool {
|
||||
if !m.errh(err) {
|
||||
return false
|
||||
}
|
||||
|
||||
if ne, ok := err.(net.Error); ok {
|
||||
return ne.Temporary()
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
type muxListener struct {
|
||||
net.Listener
|
||||
cch chan net.Conn
|
||||
}
|
||||
|
||||
func (l muxListener) Accept() (c net.Conn, err error) {
|
||||
c, ok := <-l.cch
|
||||
if !ok {
|
||||
return nil, ErrListenerClosed
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
type MuxConn struct {
|
||||
net.Conn
|
||||
prv *bytes.Buffer
|
||||
nxt *bytes.Buffer
|
||||
}
|
||||
|
||||
func newMuxConn(c net.Conn) *MuxConn {
|
||||
return &MuxConn{
|
||||
Conn: c,
|
||||
prv: &bytes.Buffer{},
|
||||
nxt: &bytes.Buffer{},
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MuxConn) Read(b []byte) (n int, err error) {
|
||||
if n, err = m.prv.Read(b); err == nil {
|
||||
return
|
||||
}
|
||||
|
||||
n, err = m.Conn.Read(b)
|
||||
return
|
||||
}
|
||||
|
||||
func (m *MuxConn) sniffer() io.Reader {
|
||||
return io.MultiReader(io.TeeReader(m.prv, m.nxt), io.TeeReader(m.Conn, m.nxt))
|
||||
}
|
||||
|
||||
func (m *MuxConn) reset() {
|
||||
if m.nxt.Len() == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
if m.prv.Len() != 0 {
|
||||
io.Copy(m.nxt, m.prv)
|
||||
}
|
||||
|
||||
m.prv, m.nxt = m.nxt, m.prv
|
||||
m.nxt.Reset()
|
||||
}
|
||||
175
cmux_test.go
Normal file
175
cmux_test.go
Normal file
@ -0,0 +1,175 @@
|
||||
package cmux
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/rpc"
|
||||
"testing"
|
||||
|
||||
"github.com/bradfitz/http2"
|
||||
)
|
||||
|
||||
const (
|
||||
testHTTP1Resp = "http1"
|
||||
rpcVal = 1234
|
||||
)
|
||||
|
||||
var testPort = 5125
|
||||
|
||||
func testAddr() string {
|
||||
testPort++
|
||||
return fmt.Sprintf("127.0.0.1:%d", testPort)
|
||||
}
|
||||
|
||||
func testListener(t *testing.T) (net.Listener, string) {
|
||||
addr := testAddr()
|
||||
l, err := net.Listen("tcp", addr)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return l, addr
|
||||
}
|
||||
|
||||
type testHTTP1Handler struct{}
|
||||
|
||||
func (h *testHTTP1Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
fmt.Fprintf(w, testHTTP1Resp)
|
||||
}
|
||||
|
||||
func runTestHTTPServer(l net.Listener, withHTTP2 bool) {
|
||||
s := &http.Server{
|
||||
Handler: &testHTTP1Handler{},
|
||||
}
|
||||
if withHTTP2 {
|
||||
http2.ConfigureServer(s, &http2.Server{})
|
||||
}
|
||||
s.Serve(l)
|
||||
}
|
||||
|
||||
func runTestHTTP1Client(t *testing.T, addr string) {
|
||||
r, err := http.Get("http://" + addr)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
defer r.Body.Close()
|
||||
b, err := ioutil.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
|
||||
if string(b) != testHTTP1Resp {
|
||||
t.Errorf("invalid response: want=%s got=%s", testHTTP1Resp, b)
|
||||
}
|
||||
}
|
||||
|
||||
type TestRPCRcvr struct{}
|
||||
|
||||
func (r TestRPCRcvr) Test(i int, j *int) error {
|
||||
*j = i
|
||||
return nil
|
||||
}
|
||||
|
||||
func runTestRPCServer(l net.Listener) {
|
||||
s := rpc.NewServer()
|
||||
s.Register(TestRPCRcvr{})
|
||||
|
||||
for {
|
||||
c, err := l.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
s.ServeConn(c)
|
||||
}
|
||||
}
|
||||
|
||||
func runTestRPCClient(t *testing.T, addr string) {
|
||||
c, err := rpc.Dial("tcp", addr)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
|
||||
var num int
|
||||
if err := c.Call("TestRPCRcvr.Test", rpcVal, &num); err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
}
|
||||
|
||||
if num != rpcVal {
|
||||
t.Errorf("wrong rpc response: want=%d got=%v", rpcVal, num)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAny(t *testing.T) {
|
||||
l, addr := testListener(t)
|
||||
defer l.Close()
|
||||
|
||||
muxl := New(l)
|
||||
httpl := muxl.Match(Any())
|
||||
|
||||
go runTestHTTPServer(httpl, false)
|
||||
go muxl.Serve()
|
||||
|
||||
r, err := http.Get("http://" + addr)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
defer r.Body.Close()
|
||||
b, err := ioutil.ReadAll(r.Body)
|
||||
if string(b) != testHTTP1Resp {
|
||||
t.Errorf("invalid response: want=%s got=%s", testHTTP1Resp, b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPGoRPC(t *testing.T) {
|
||||
l, addr := testListener(t)
|
||||
defer l.Close()
|
||||
|
||||
muxl := New(l)
|
||||
httpl := muxl.Match(HTTP2(), HTTP1Fast())
|
||||
rpcl := muxl.Match(Any())
|
||||
|
||||
go runTestHTTPServer(httpl, true)
|
||||
go runTestRPCServer(rpcl)
|
||||
go muxl.Serve()
|
||||
|
||||
runTestHTTP1Client(t, addr)
|
||||
runTestRPCClient(t, addr)
|
||||
}
|
||||
|
||||
func TestErrorHandler(t *testing.T) {
|
||||
l, addr := testListener(t)
|
||||
defer l.Close()
|
||||
|
||||
muxl := New(l)
|
||||
httpl := muxl.Match(HTTP2(), HTTP1Fast())
|
||||
|
||||
go runTestHTTPServer(httpl, true)
|
||||
go muxl.Serve()
|
||||
|
||||
firstErr := true
|
||||
muxl.HandleError(func(err error) bool {
|
||||
if !firstErr {
|
||||
return true
|
||||
}
|
||||
if _, ok := err.(ErrNotMatched); !ok {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
firstErr = false
|
||||
return true
|
||||
})
|
||||
|
||||
c, err := rpc.Dial("tcp", addr)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var num int
|
||||
if err := c.Call("TestRPCRcvr.Test", rpcVal, &num); err == nil {
|
||||
t.Error("rpc got a response")
|
||||
}
|
||||
}
|
||||
81
example_test.go
Normal file
81
example_test.go
Normal file
@ -0,0 +1,81 @@
|
||||
package cmux_test
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/rpc"
|
||||
|
||||
"google.golang.org/grpc"
|
||||
|
||||
"golang.org/x/net/context"
|
||||
|
||||
grpchello "github.com/grpc/grpc-common/go/helloworld"
|
||||
"github.com/soheilhy/cmux"
|
||||
)
|
||||
|
||||
type exampleHTTPHandler struct{}
|
||||
|
||||
func (h *exampleHTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
fmt.Fprintf(w, "example http response")
|
||||
}
|
||||
|
||||
func serveHTTP(l net.Listener) {
|
||||
s := &http.Server{
|
||||
Handler: &exampleHTTPHandler{},
|
||||
}
|
||||
s.Serve(l)
|
||||
}
|
||||
|
||||
type ExampleRPCRcvr struct{}
|
||||
|
||||
func (r *ExampleRPCRcvr) Cube(i int, j *int) error {
|
||||
*j = i * i
|
||||
return nil
|
||||
}
|
||||
|
||||
func serveRPC(l net.Listener) {
|
||||
s := rpc.NewServer()
|
||||
s.Register(&ExampleRPCRcvr{})
|
||||
s.Accept(l)
|
||||
}
|
||||
|
||||
type grpcServer struct{}
|
||||
|
||||
func (s *grpcServer) SayHello(ctx context.Context, in *grpchello.HelloRequest) (
|
||||
*grpchello.HelloReply, error) {
|
||||
|
||||
return &grpchello.HelloReply{Message: "Hello " + in.Name + " from cmux"}, nil
|
||||
}
|
||||
|
||||
func serveGRPC(l net.Listener) {
|
||||
grpcs := grpc.NewServer()
|
||||
grpchello.RegisterGreeterServer(grpcs, &grpcServer{})
|
||||
grpcs.Serve(l)
|
||||
}
|
||||
|
||||
func Example() {
|
||||
l, err := net.Listen("tcp", "127.0.0.1:50051")
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
m := cmux.New(l)
|
||||
|
||||
// We first match the connection against HTTP2 fields. If matched, the
|
||||
// connection will be sent through the "grpcl" listener.
|
||||
grpcl := m.Match(cmux.HTTP2HeaderField("content-type", "application/grpc"))
|
||||
// Otherwise, we match it againts HTTP1 methods and HTTP2. If matched by
|
||||
// any of them, it is sent through the "httpl" listener.
|
||||
httpl := m.Match(cmux.HTTP1Fast(), cmux.HTTP2())
|
||||
// If not matched by HTTP, we assume it is an RPC connection.
|
||||
rpcl := m.Match(cmux.Any())
|
||||
|
||||
// Then we used the muxed listeners.
|
||||
go serveGRPC(grpcl)
|
||||
go serveHTTP(httpl)
|
||||
go serveRPC(rpcl)
|
||||
|
||||
m.Serve()
|
||||
}
|
||||
173
patricia.go
Normal file
173
patricia.go
Normal file
@ -0,0 +1,173 @@
|
||||
package cmux
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
)
|
||||
|
||||
// patriciaTree is a simple patricia tree that handles []byte instead of string
|
||||
// and cannot be changed after instantiation.
|
||||
type patriciaTree struct {
|
||||
root *ptNode
|
||||
}
|
||||
|
||||
func newPatriciaTree(b ...[]byte) *patriciaTree {
|
||||
return &patriciaTree{
|
||||
root: newNode(b),
|
||||
}
|
||||
}
|
||||
|
||||
func newPatriciaTreeString(strs ...string) *patriciaTree {
|
||||
b := make([][]byte, len(strs))
|
||||
for i, s := range strs {
|
||||
b[i] = []byte(s)
|
||||
}
|
||||
return &patriciaTree{
|
||||
root: newNode(b),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *patriciaTree) matchPrefix(r io.Reader) bool {
|
||||
return t.root.match(r, true)
|
||||
}
|
||||
|
||||
func (t *patriciaTree) match(r io.Reader) bool {
|
||||
return t.root.match(r, false)
|
||||
}
|
||||
|
||||
type ptNode struct {
|
||||
prefix []byte
|
||||
next map[byte]*ptNode
|
||||
terminal bool
|
||||
}
|
||||
|
||||
func newNode(strs [][]byte) *ptNode {
|
||||
if len(strs) == 0 {
|
||||
return &ptNode{
|
||||
prefix: []byte{},
|
||||
terminal: true,
|
||||
}
|
||||
}
|
||||
|
||||
if len(strs) == 1 {
|
||||
return &ptNode{
|
||||
prefix: strs[0],
|
||||
terminal: true,
|
||||
}
|
||||
}
|
||||
|
||||
p, strs := splitPrefix(strs)
|
||||
n := &ptNode{
|
||||
prefix: p,
|
||||
}
|
||||
|
||||
nexts := make(map[byte][][]byte)
|
||||
for _, s := range strs {
|
||||
if len(s) == 0 {
|
||||
n.terminal = true
|
||||
continue
|
||||
}
|
||||
nexts[s[0]] = append(nexts[s[0]], s[1:])
|
||||
}
|
||||
|
||||
n.next = make(map[byte]*ptNode)
|
||||
for first, rests := range nexts {
|
||||
n.next[first] = newNode(rests)
|
||||
}
|
||||
|
||||
return n
|
||||
}
|
||||
|
||||
func splitPrefix(bss [][]byte) (prefix []byte, rest [][]byte) {
|
||||
if len(bss) == 0 || len(bss[0]) == 0 {
|
||||
return prefix, bss
|
||||
}
|
||||
|
||||
if len(bss) == 1 {
|
||||
return bss[0], [][]byte{[]byte{}}
|
||||
}
|
||||
|
||||
for i := 0; ; i++ {
|
||||
var cur byte
|
||||
eq := true
|
||||
for j, b := range bss {
|
||||
if len(b) <= i {
|
||||
eq = false
|
||||
break
|
||||
}
|
||||
|
||||
if j == 0 {
|
||||
cur = b[i]
|
||||
continue
|
||||
}
|
||||
|
||||
if cur != b[i] {
|
||||
eq = false
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if !eq {
|
||||
break
|
||||
}
|
||||
|
||||
prefix = append(prefix, cur)
|
||||
}
|
||||
|
||||
rest = make([][]byte, 0, len(bss))
|
||||
for _, b := range bss {
|
||||
rest = append(rest, b[len(prefix):])
|
||||
}
|
||||
|
||||
return prefix, rest
|
||||
}
|
||||
|
||||
func readBytes(r io.Reader, n int) (b []byte, err error) {
|
||||
b = make([]byte, n)
|
||||
o := 0
|
||||
for o < n {
|
||||
nr, err := r.Read(b[o:])
|
||||
if err != nil && err != io.EOF {
|
||||
return b, err
|
||||
}
|
||||
|
||||
o += nr
|
||||
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
}
|
||||
return b[:o], nil
|
||||
}
|
||||
|
||||
func (n *ptNode) match(r io.Reader, prefix bool) bool {
|
||||
if l := len(n.prefix); l > 0 {
|
||||
b, err := readBytes(r, l)
|
||||
if err != nil || len(b) != l || !bytes.Equal(b, n.prefix) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
if prefix && n.terminal {
|
||||
return true
|
||||
}
|
||||
|
||||
b := make([]byte, 1)
|
||||
for {
|
||||
nr, err := r.Read(b)
|
||||
if nr != 0 {
|
||||
break
|
||||
}
|
||||
|
||||
if err == io.EOF {
|
||||
return n.terminal
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
nextN, ok := n.next[b[0]]
|
||||
return ok && nextN.match(r, prefix)
|
||||
}
|
||||
35
patricia_test.go
Normal file
35
patricia_test.go
Normal file
@ -0,0 +1,35 @@
|
||||
package cmux
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func testPTree(t *testing.T, strs ...string) {
|
||||
pt := newPatriciaTreeString(strs...)
|
||||
for _, s := range strs {
|
||||
if !pt.match(strings.NewReader(s)) {
|
||||
t.Errorf("%s is not matched by %s", s, s)
|
||||
}
|
||||
|
||||
if !pt.matchPrefix(strings.NewReader(s + s)) {
|
||||
t.Errorf("%s is not matched as a prefix by %s", s+s, s)
|
||||
}
|
||||
|
||||
if pt.match(strings.NewReader(s + s)) {
|
||||
t.Errorf("%s matches %s", s+s, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPatriciaOnePrefix(t *testing.T) {
|
||||
testPTree(t, "prefix")
|
||||
}
|
||||
|
||||
func TestPatriciaNonOverlapping(t *testing.T) {
|
||||
testPTree(t, "foo", "bar", "dummy")
|
||||
}
|
||||
|
||||
func TestPatriciaOverlapping(t *testing.T) {
|
||||
testPTree(t, "foo", "far", "farther", "boo", "bar")
|
||||
}
|
||||
156
selectors.go
Normal file
156
selectors.go
Normal file
@ -0,0 +1,156 @@
|
||||
package cmux
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/bradfitz/http2"
|
||||
"github.com/bradfitz/http2/hpack"
|
||||
)
|
||||
|
||||
// Any is a Matcher that matches any connection.
|
||||
func Any() Matcher {
|
||||
return func(r io.Reader) bool { return true }
|
||||
}
|
||||
|
||||
// PrefixMatcher returns a matcher that matches a connection if it
|
||||
// starts with any of the strings in strs.
|
||||
func PrefixMatcher(strs ...string) Matcher {
|
||||
pt := newPatriciaTreeString(strs...)
|
||||
return func(r io.Reader) bool {
|
||||
return pt.matchPrefix(r)
|
||||
}
|
||||
}
|
||||
|
||||
var defaultHTTPMethods = []string{
|
||||
"OPTIONS",
|
||||
"GET",
|
||||
"HEAD",
|
||||
"POST",
|
||||
"PUT",
|
||||
"DELETE",
|
||||
"TRACE",
|
||||
"CONNECT",
|
||||
}
|
||||
|
||||
// HTTP1Fast only matches the methods in the HTTP request.
|
||||
//
|
||||
// This matcher is very optimistic: if it returns true, it does not mean that
|
||||
// the request is a valid HTTP response. If you want a correct but slower HTTP1
|
||||
// matcher, use HTTP1 instead.
|
||||
func HTTP1Fast(extMethods ...string) Matcher {
|
||||
return PrefixMatcher(append(defaultHTTPMethods, extMethods...)...)
|
||||
}
|
||||
|
||||
const (
|
||||
maxHTTPRead = 4096
|
||||
)
|
||||
|
||||
// HTTP1 parses the first line or upto 4096 bytes of the request to see if
|
||||
// the conection contains an HTTP request.
|
||||
func HTTP1() Matcher {
|
||||
return func(r io.Reader) bool {
|
||||
br := bufio.NewReader(&io.LimitedReader{R: r, N: maxHTTPRead})
|
||||
l, part, err := br.ReadLine()
|
||||
if err != nil || part {
|
||||
return false
|
||||
}
|
||||
|
||||
_, _, proto, ok := parseRequestLine(string(l))
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
|
||||
v, _, ok := http.ParseHTTPVersion(proto)
|
||||
return ok && v == 1
|
||||
}
|
||||
}
|
||||
|
||||
// grabbed from net/http.
|
||||
func parseRequestLine(line string) (method, uri, proto string, ok bool) {
|
||||
s1 := strings.Index(line, " ")
|
||||
s2 := strings.Index(line[s1+1:], " ")
|
||||
if s1 < 0 || s2 < 0 {
|
||||
return
|
||||
}
|
||||
s2 += s1 + 1
|
||||
return line[:s1], line[s1+1 : s2], line[s2+1:], true
|
||||
}
|
||||
|
||||
var (
|
||||
http2Preface = []byte(http2.ClientPreface)
|
||||
)
|
||||
|
||||
// HTTP2 parses the frame header of the first frame to detect whether the
|
||||
// connection is an HTTP2 connection.
|
||||
func HTTP2() Matcher {
|
||||
return func(r io.Reader) bool {
|
||||
return hasHTTP2Preface(r)
|
||||
}
|
||||
}
|
||||
|
||||
// HTTP1HeaderField returns a matcher matching the header fields of the first
|
||||
// request of an HTTP 1 connection.
|
||||
func HTTP1HeaderField(name, value string) Matcher {
|
||||
return func(r io.Reader) bool {
|
||||
return matchHTTP1Field(r, name, value)
|
||||
}
|
||||
}
|
||||
|
||||
// HTTP2HeaderField resturns a matcher matching the header fields of the first
|
||||
// headers frame.
|
||||
func HTTP2HeaderField(name, value string) Matcher {
|
||||
return func(r io.Reader) bool {
|
||||
return matchHTTP2Field(r, name, value)
|
||||
}
|
||||
}
|
||||
|
||||
func hasHTTP2Preface(r io.Reader) (ok bool) {
|
||||
b := make([]byte, len(http2Preface))
|
||||
n, err := r.Read(b)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
b = b[:n]
|
||||
return bytes.Equal(b, http2Preface)
|
||||
}
|
||||
|
||||
func matchHTTP1Field(r io.Reader, name, value string) (matched bool) {
|
||||
return
|
||||
}
|
||||
|
||||
func matchHTTP2Field(r io.Reader, name, value string) (matched bool) {
|
||||
if !hasHTTP2Preface(r) {
|
||||
return false
|
||||
}
|
||||
|
||||
framer := http2.NewFramer(ioutil.Discard, r)
|
||||
hdec := hpack.NewDecoder(uint32(4<<10), func(hf hpack.HeaderField) {
|
||||
if hf.Name == name && hf.Value == value {
|
||||
matched = true
|
||||
}
|
||||
})
|
||||
for {
|
||||
f, err := framer.ReadFrame()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
switch f := f.(type) {
|
||||
case *http2.HeadersFrame:
|
||||
hdec.Write(f.HeaderBlockFragment())
|
||||
if matched {
|
||||
return true
|
||||
}
|
||||
|
||||
if f.FrameHeader.Flags&http2.FlagHeadersEndHeaders != 0 {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user