mirror of
https://github.com/sot-tech/mochi.git
synced 2026-07-09 10:08:11 -07:00
initial middleware refactor
This commit is contained in:
Jimmy Zelinskie
@@ -0,0 +1,33 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package http
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"gopkg.in/yaml.v2"
|
||||
)
|
||||
|
||||
type httpConfig struct {
|
||||
Addr string `yaml:"addr"`
|
||||
RequestTimeout time.Duration `yaml:"requestTimeout"`
|
||||
ReadTimeout time.Duration `yaml:"readTimeout"`
|
||||
WriteTimeout time.Duration `yaml:"writeTimeout"`
|
||||
}
|
||||
|
||||
func newHTTPConfig(srvcfg interface{}) (*httpConfig, error) {
|
||||
bytes, err := yaml.Marshal(srvcfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var cfg httpConfig
|
||||
err = yaml.Unmarshal(bytes, &cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &cfg, nil
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
// Package query implements a simple, fast URL parser designed to be used to
|
||||
// parse parameters sent from BitTorrent clients. The last value of a key wins,
|
||||
// except for they key "info_hash".
|
||||
package query
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/chihaya/chihaya"
|
||||
"github.com/chihaya/chihaya/pkg/event"
|
||||
)
|
||||
|
||||
// ErrKeyNotFound is returned when a provided key has no value associated with
|
||||
// it.
|
||||
var ErrKeyNotFound = errors.New("query: value for the provided key does not exist")
|
||||
|
||||
// Query represents a parsed URL.Query.
|
||||
type Query struct {
|
||||
query string
|
||||
infohashes []string
|
||||
params map[string]string
|
||||
}
|
||||
|
||||
// New parses a raw URL query.
|
||||
func New(query string) (*Query, error) {
|
||||
var (
|
||||
keyStart, keyEnd int
|
||||
valStart, valEnd int
|
||||
firstInfohash string
|
||||
|
||||
onKey = true
|
||||
hasInfohash = false
|
||||
|
||||
q = &Query{
|
||||
query: query,
|
||||
infohashes: nil,
|
||||
params: make(map[string]string),
|
||||
}
|
||||
)
|
||||
|
||||
for i, length := 0, len(query); i < length; i++ {
|
||||
separator := query[i] == '&' || query[i] == ';' || query[i] == '?'
|
||||
last := i == length-1
|
||||
|
||||
if separator || last {
|
||||
if onKey && !last {
|
||||
keyStart = i + 1
|
||||
continue
|
||||
}
|
||||
|
||||
if last && !separator && !onKey {
|
||||
valEnd = i
|
||||
}
|
||||
|
||||
keyStr, err := url.QueryUnescape(query[keyStart : keyEnd+1])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var valStr string
|
||||
|
||||
if valEnd > 0 {
|
||||
valStr, err = url.QueryUnescape(query[valStart : valEnd+1])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
q.params[strings.ToLower(keyStr)] = valStr
|
||||
|
||||
if keyStr == "info_hash" {
|
||||
if hasInfohash {
|
||||
// Multiple infohashes
|
||||
if q.infohashes == nil {
|
||||
q.infohashes = []string{firstInfohash}
|
||||
}
|
||||
q.infohashes = append(q.infohashes, valStr)
|
||||
} else {
|
||||
firstInfohash = valStr
|
||||
hasInfohash = true
|
||||
}
|
||||
}
|
||||
|
||||
valEnd = 0
|
||||
onKey = true
|
||||
keyStart = i + 1
|
||||
|
||||
} else if query[i] == '=' {
|
||||
onKey = false
|
||||
valStart = i + 1
|
||||
valEnd = 0
|
||||
} else if onKey {
|
||||
keyEnd = i
|
||||
} else {
|
||||
valEnd = i
|
||||
}
|
||||
}
|
||||
|
||||
return q, nil
|
||||
}
|
||||
|
||||
// Infohashes returns a list of requested infohashes.
|
||||
func (q *Query) Infohashes() ([]string, error) {
|
||||
if q.infohashes == nil {
|
||||
infohash, err := q.String("info_hash")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []string{infohash}, nil
|
||||
}
|
||||
return q.infohashes, nil
|
||||
}
|
||||
|
||||
// String returns a string parsed from a query. Every key can be returned as a
|
||||
// string because they are encoded in the URL as strings.
|
||||
func (q *Query) String(key string) (string, error) {
|
||||
val, exists := q.params[key]
|
||||
if !exists {
|
||||
return "", ErrKeyNotFound
|
||||
}
|
||||
return val, nil
|
||||
}
|
||||
|
||||
// Uint64 returns a uint parsed from a query. After being called, it is safe to
|
||||
// cast the uint64 to your desired length.
|
||||
func (q *Query) Uint64(key string) (uint64, error) {
|
||||
str, exists := q.params[key]
|
||||
if !exists {
|
||||
return 0, ErrKeyNotFound
|
||||
}
|
||||
|
||||
val, err := strconv.ParseUint(str, 10, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return val, nil
|
||||
}
|
||||
|
||||
// AnnounceRequest generates an chihaya.AnnounceRequest with the parameters
|
||||
// provided by a query.
|
||||
func (q *Query) AnnounceRequest() (chihaya.AnnounceRequest, error) {
|
||||
request := make(chihaya.AnnounceRequest)
|
||||
|
||||
request["query"] = q.query
|
||||
|
||||
eventStr, err := q.String("event")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: event")
|
||||
}
|
||||
request["event"], err = event.New(eventStr)
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to provide valid client event")
|
||||
}
|
||||
|
||||
compactStr, err := q.String("compact")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: compact")
|
||||
}
|
||||
request["compact"] = compactStr != "0"
|
||||
|
||||
request["info_hash"], err = q.String("info_hash")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: info_hash")
|
||||
}
|
||||
|
||||
request["peer_id"], err = q.String("peer_id")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: peer_id")
|
||||
}
|
||||
|
||||
request["left"], err = q.Uint64("left")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: left")
|
||||
}
|
||||
|
||||
request["downloaded"], err = q.Uint64("downloaded")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: downloaded")
|
||||
}
|
||||
|
||||
request["uploaded"], err = q.Uint64("uploaded")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: uploaded")
|
||||
}
|
||||
|
||||
request["numwant"], err = q.String("numwant")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: numwant")
|
||||
}
|
||||
|
||||
request["ip"], err = q.Uint64("port")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: port")
|
||||
}
|
||||
|
||||
request["port"], err = q.Uint64("port")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: port")
|
||||
}
|
||||
|
||||
request["ip"], err = q.String("ip")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: ip")
|
||||
}
|
||||
|
||||
request["ipv4"], err = q.String("ipv4")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: ipv4")
|
||||
}
|
||||
|
||||
request["ipv6"], err = q.String("ipv6")
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: ipv6")
|
||||
}
|
||||
|
||||
return request, nil
|
||||
}
|
||||
|
||||
// ScrapeRequest generates an chihaya.ScrapeRequeset with the parameters
|
||||
// provided by a query.
|
||||
func (q *Query) ScrapeRequest() (chihaya.ScrapeRequest, error) {
|
||||
request := make(chihaya.ScrapeRequest)
|
||||
|
||||
var err error
|
||||
request["info_hash"], err = q.Infohashes()
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to parse parameter: info_hash")
|
||||
}
|
||||
|
||||
return request, nil
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package http
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
|
||||
"github.com/julienschmidt/httprouter"
|
||||
"github.com/tylerb/graceful"
|
||||
|
||||
"github.com/chihaya/chihaya/config"
|
||||
"github.com/chihaya/chihaya/server"
|
||||
"github.com/chihaya/chihaya/server/http/query"
|
||||
"github.com/chihaya/chihaya/tracker"
|
||||
)
|
||||
|
||||
func init() {
|
||||
server.Register("http", constructor)
|
||||
}
|
||||
|
||||
func constructor(srvcfg *config.ServerConfig, tkr *tracker.Tracker) (server.Server, error) {
|
||||
cfg, err := newHTTPConfig(srvcfg)
|
||||
if err != nil {
|
||||
return nil, errors.New("http: invalid config: " + err.Error())
|
||||
}
|
||||
|
||||
return &httpServer{
|
||||
cfg: cfg,
|
||||
tkr: tkr,
|
||||
}, nil
|
||||
}
|
||||
|
||||
type httpServer struct {
|
||||
cfg *httpConfig
|
||||
tkr *tracker.Tracker
|
||||
grace *graceful.Server
|
||||
stopping bool
|
||||
}
|
||||
|
||||
func (s *httpServer) Start() {
|
||||
s.grace = &graceful.Server{
|
||||
Server: &http.Server{
|
||||
Addr: s.cfg.Addr,
|
||||
Handler: s.routes(),
|
||||
ReadTimeout: s.cfg.ReadTimeout,
|
||||
WriteTimeout: s.cfg.WriteTimeout,
|
||||
},
|
||||
Timeout: s.cfg.RequestTimeout,
|
||||
NoSignalHandling: true,
|
||||
ShutdownInitiated: func() { s.stopping = true },
|
||||
ConnState: func(conn net.Conn, state http.ConnState) {
|
||||
switch state {
|
||||
case http.StateNew:
|
||||
//stats.RecordEvent(stats.AcceptedConnection)
|
||||
|
||||
case http.StateClosed:
|
||||
//stats.RecordEvent(stats.ClosedConnection)
|
||||
|
||||
case http.StateHijacked:
|
||||
panic("http: connection impossibly hijacked")
|
||||
|
||||
// Ignore the following cases.
|
||||
case http.StateActive, http.StateIdle:
|
||||
|
||||
default:
|
||||
panic("http: connection transitioned to unknown state")
|
||||
}
|
||||
},
|
||||
}
|
||||
s.grace.SetKeepAlivesEnabled(false)
|
||||
|
||||
if err := s.grace.ListenAndServe(); err != nil {
|
||||
if opErr, ok := err.(*net.OpError); !ok || (ok && opErr.Op != "accept") {
|
||||
log.Printf("Failed to gracefully run HTTP server: %s", err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *httpServer) Stop() {
|
||||
if !s.stopping {
|
||||
s.grace.Stop(s.grace.Timeout)
|
||||
}
|
||||
|
||||
s.grace = nil
|
||||
s.stopping = false
|
||||
}
|
||||
|
||||
func (s *httpServer) routes() *httprouter.Router {
|
||||
r := httprouter.New()
|
||||
r.GET("/announce", s.serveAnnounce)
|
||||
r.GET("/scrape", s.serveScrape)
|
||||
return r
|
||||
}
|
||||
|
||||
func (s *httpServer) serveAnnounce(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
|
||||
writer := &writer{w}
|
||||
|
||||
q, err := query.New(r.URL.RawQuery)
|
||||
if err != nil {
|
||||
writer.writeError(err)
|
||||
return
|
||||
}
|
||||
|
||||
req, err := q.AnnounceRequest()
|
||||
if err != nil {
|
||||
writer.writeError(err)
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := s.tkr.HandleAnnounce(req)
|
||||
if err != nil {
|
||||
writer.writeError(err)
|
||||
return
|
||||
}
|
||||
|
||||
writer.writeAnnounceResponse(resp)
|
||||
}
|
||||
|
||||
func (s *httpServer) serveScrape(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
|
||||
writer := &writer{w}
|
||||
|
||||
q, err := query.New(r.URL.RawQuery)
|
||||
if err != nil {
|
||||
writer.writeError(err)
|
||||
return
|
||||
}
|
||||
|
||||
req, err := q.ScrapeRequest()
|
||||
if err != nil {
|
||||
writer.writeError(err)
|
||||
return
|
||||
}
|
||||
|
||||
resp, err := s.tkr.HandleScrape(req)
|
||||
if err != nil {
|
||||
writer.writeError(err)
|
||||
return
|
||||
}
|
||||
|
||||
writer.writeScrapeResponse(resp)
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package http
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/chihaya/chihaya"
|
||||
"github.com/chihaya/chihaya/pkg/bencode"
|
||||
)
|
||||
|
||||
type writer struct{ http.ResponseWriter }
|
||||
|
||||
func (w *writer) writeError(err error) error {
|
||||
return bencode.NewEncoder(w).Encode(bencode.Dict{
|
||||
"failure reason": err.Error(),
|
||||
})
|
||||
}
|
||||
|
||||
func (w *writer) writeAnnounceResponse(resp *chihaya.AnnounceResponse) error {
|
||||
bdict := bencode.Dict{
|
||||
"complete": resp.Complete,
|
||||
"incomplete": resp.Incomplete,
|
||||
"interval": resp.Interval,
|
||||
"min interval": resp.MinInterval,
|
||||
}
|
||||
|
||||
// Add the peers to the dictionary in the compact format.
|
||||
if resp.Compact {
|
||||
var IPv4CompactDict, IPv6CompactDict []byte
|
||||
|
||||
// Add the IPv4 peers to the dictionary.
|
||||
for _, peer := range resp.IPv4Peers {
|
||||
IPv4CompactDict = append(IPv4CompactDict, compact(peer)...)
|
||||
}
|
||||
if len(IPv4CompactDict) > 0 {
|
||||
bdict["peers"] = IPv4CompactDict
|
||||
}
|
||||
|
||||
// Add the IPv6 peers to the dictionary.
|
||||
for _, peer := range resp.IPv6Peers {
|
||||
IPv6CompactDict = append(IPv6CompactDict, compact(peer)...)
|
||||
}
|
||||
if len(IPv6CompactDict) > 0 {
|
||||
bdict["peers6"] = IPv6CompactDict
|
||||
}
|
||||
|
||||
return bencode.NewEncoder(w).Encode(bdict)
|
||||
}
|
||||
|
||||
// Add the peers to the dictionary.
|
||||
var peers []bencode.Dict
|
||||
for _, peer := range resp.IPv4Peers {
|
||||
peers = append(peers, dict(peer))
|
||||
}
|
||||
for _, peer := range resp.IPv6Peers {
|
||||
peers = append(peers, dict(peer))
|
||||
}
|
||||
bdict["peers"] = peers
|
||||
|
||||
return bencode.NewEncoder(w).Encode(bdict)
|
||||
}
|
||||
|
||||
func (w *writer) writeScrapeResponse(resp *chihaya.ScrapeResponse) error {
|
||||
filesDict := bencode.NewDict()
|
||||
for infohash, scrape := range resp.Files {
|
||||
filesDict[infohash] = bencode.Dict{
|
||||
"complete": scrape.Complete,
|
||||
"incomplete": scrape.Incomplete,
|
||||
"downloaded": scrape.Downloaded,
|
||||
}
|
||||
}
|
||||
|
||||
return bencode.NewEncoder(w).Encode(bencode.Dict{
|
||||
"files": filesDict,
|
||||
})
|
||||
}
|
||||
|
||||
func compact(peer chihaya.Peer) (buf []byte) {
|
||||
buf = []byte(peer.IP)
|
||||
buf = append(buf, byte(peer.Port>>8))
|
||||
buf = append(buf, byte(peer.Port&0xff))
|
||||
return
|
||||
}
|
||||
|
||||
func dict(peer chihaya.Peer) bencode.Dict {
|
||||
return bencode.Dict{
|
||||
"peer id": peer.ID,
|
||||
"ip": peer.IP.String(),
|
||||
"port": peer.Port,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package http
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestWriteError(t *testing.T) {
|
||||
var table = []struct {
|
||||
reason, expected string
|
||||
}{
|
||||
{"hello world", "d14:failure reason11:hello worlde"},
|
||||
{"what's up", "d14:failure reason9:what's upe"},
|
||||
}
|
||||
|
||||
for _, tt := range table {
|
||||
r := httptest.NewRecorder()
|
||||
w := &writer{r}
|
||||
err := w.writeError(errors.New(tt.reason))
|
||||
assert.Nil(t, err, "writeError should not fail with test input")
|
||||
assert.Equal(t, r.Body.String(), tt.expected, "writer should write the expected value")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/chihaya/chihaya/config"
|
||||
"github.com/chihaya/chihaya/tracker"
|
||||
)
|
||||
|
||||
// Pool represents a running pool of servers.
|
||||
type Pool struct {
|
||||
servers []Server
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// StartPool creates a new pool of servers specified by the provided config and
|
||||
// runs them.
|
||||
func StartPool(cfgs []config.ServerConfig, tkr *tracker.Tracker) (*Pool, error) {
|
||||
var servers []Server
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for _, cfg := range cfgs {
|
||||
srv, err := New(&cfg, tkr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
go func(srv Server) {
|
||||
defer wg.Done()
|
||||
srv.Start()
|
||||
}(srv)
|
||||
|
||||
servers = append(servers, srv)
|
||||
}
|
||||
|
||||
return &Pool{
|
||||
servers: servers,
|
||||
wg: wg,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Stop safely shuts down a pool of servers.
|
||||
func (p *Pool) Stop() {
|
||||
for _, srv := range p.servers {
|
||||
srv.Stop()
|
||||
}
|
||||
p.wg.Wait()
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
// Package server implements an abstraction over servers meant to be run .
|
||||
// alongside a tracker.
|
||||
//
|
||||
// Servers may be implementations of different transport protocols or have their
|
||||
// own custom behavior.
|
||||
package server
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/chihaya/chihaya/config"
|
||||
"github.com/chihaya/chihaya/tracker"
|
||||
)
|
||||
|
||||
var constructors = make(map[string]Constructor)
|
||||
|
||||
// Constructor is a function that creates a new Server.
|
||||
type Constructor func(*config.ServerConfig, *tracker.Tracker) (Server, error)
|
||||
|
||||
// Register makes a Constructor available by the provided name.
|
||||
//
|
||||
// If this function is called twice with the same name or if the Constructor is
|
||||
// nil, it panics.
|
||||
func Register(name string, con Constructor) {
|
||||
if con == nil {
|
||||
panic("server: could not register nil Constructor")
|
||||
}
|
||||
if _, dup := constructors[name]; dup {
|
||||
panic("server: could not register duplicate Constructor: " + name)
|
||||
}
|
||||
constructors[name] = con
|
||||
}
|
||||
|
||||
// New creates a Server specified by a configuration.
|
||||
func New(cfg *config.ServerConfig, tkr *tracker.Tracker) (Server, error) {
|
||||
con, ok := constructors[cfg.Name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf(
|
||||
"server: unknown Constructor %q (forgotten import?)",
|
||||
cfg.Name,
|
||||
)
|
||||
}
|
||||
return con(cfg, tkr)
|
||||
}
|
||||
|
||||
// Server represents one instance of a server accessing the tracker.
|
||||
type Server interface {
|
||||
Start()
|
||||
Stop()
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package store
|
||||
|
||||
import "fmt"
|
||||
|
||||
var clientStoreDrivers = make(map[string]ClientStoreDriver)
|
||||
|
||||
// ClientStore represents an interface for manipulating clientIDs.
|
||||
type ClientStore interface {
|
||||
CreateClient(clientID string) error
|
||||
FindClient(peerID string) (bool, error)
|
||||
DeleteClient(clientID string) error
|
||||
}
|
||||
|
||||
// ClientStoreDriver represents an interface for creating a handle to the
|
||||
// storage of swarms.
|
||||
type ClientStoreDriver interface {
|
||||
New(*Config) (ClientStore, error)
|
||||
}
|
||||
|
||||
// RegisterClientStoreDriver makes a driver available by the provided name.
|
||||
//
|
||||
// If this function is called twice with the same name or if the driver is nil,
|
||||
// it panics.
|
||||
func RegisterClientStoreDriver(name string, driver ClientStoreDriver) {
|
||||
if driver == nil {
|
||||
panic("store: could not register nil ClientStoreDriver")
|
||||
}
|
||||
if _, dup := clientStoreDrivers[name]; dup {
|
||||
panic("store: could not register duplicate ClientStoreDriver: " + name)
|
||||
}
|
||||
clientStoreDrivers[name] = driver
|
||||
}
|
||||
|
||||
// OpenClientStore returns a ClientStore specified by a configuration.
|
||||
func OpenClientStore(name string, cfg *Config) (ClientStore, error) {
|
||||
driver, ok := clientStoreDrivers[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf(
|
||||
"store: unknown driver %q (forgotten import?)",
|
||||
name,
|
||||
)
|
||||
}
|
||||
|
||||
return driver.New(cfg)
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package memory
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/chihaya/chihaya/pkg/clientid"
|
||||
"github.com/chihaya/chihaya/server/store"
|
||||
)
|
||||
|
||||
func init() {
|
||||
store.RegisterClientStoreDriver("memory", &clientStoreDriver{})
|
||||
}
|
||||
|
||||
type clientStoreDriver struct{}
|
||||
|
||||
func (d *clientStoreDriver) New(cfg *store.Config) (store.ClientStore, error) {
|
||||
return &clientStore{
|
||||
clientIDs: make(map[string]struct{}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
type clientStore struct {
|
||||
clientIDs map[string]struct{}
|
||||
sync.RWMutex
|
||||
}
|
||||
|
||||
var _ store.ClientStore = &clientStore{}
|
||||
|
||||
func (s *clientStore) CreateClient(clientID string) error {
|
||||
s.Lock()
|
||||
defer s.Unlock()
|
||||
|
||||
s.clientIDs[clientID] = struct{}{}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *clientStore) FindClient(peerID string) (bool, error) {
|
||||
clientID := clientid.New(peerID)
|
||||
s.Lock()
|
||||
defer s.Unlock()
|
||||
|
||||
_, ok := s.clientIDs[clientID]
|
||||
|
||||
return ok, nil
|
||||
}
|
||||
|
||||
func (s *clientStore) DeleteClient(clientID string) error {
|
||||
s.Lock()
|
||||
defer s.Unlock()
|
||||
|
||||
delete(s.clientIDs, clientID)
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,272 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package memory
|
||||
|
||||
import (
|
||||
"hash/fnv"
|
||||
"runtime"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gopkg.in/yaml.v2"
|
||||
|
||||
"github.com/chihaya/chihaya"
|
||||
"github.com/chihaya/chihaya/server/store"
|
||||
)
|
||||
|
||||
func init() {
|
||||
store.RegisterPeerStoreDriver("memory", &peerStoreDriver{})
|
||||
}
|
||||
|
||||
type peerStoreDriver struct{}
|
||||
|
||||
func (d *peerStoreDriver) New(storecfg *store.Config) (store.PeerStore, error) {
|
||||
cfg, err := newPeerStoreConfig(storecfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &peerStore{
|
||||
shards: make([]*peerShard, cfg.Shards),
|
||||
}, nil
|
||||
}
|
||||
|
||||
type peerStoreConfig struct {
|
||||
Shards int `yaml:"shards"`
|
||||
}
|
||||
|
||||
func newPeerStoreConfig(storecfg *store.Config) (*peerStoreConfig, error) {
|
||||
bytes, err := yaml.Marshal(storecfg.PeerStoreConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var cfg peerStoreConfig
|
||||
err = yaml.Unmarshal(bytes, &cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
const seedersSuffix = "-seeders"
|
||||
const leechersSuffix = "-leechers"
|
||||
|
||||
type peer struct {
|
||||
chihaya.Peer
|
||||
LastAction time.Time
|
||||
}
|
||||
|
||||
type peerShard struct {
|
||||
peers map[string]map[string]peer
|
||||
sync.RWMutex
|
||||
}
|
||||
|
||||
type peerStore struct {
|
||||
shards []*peerShard
|
||||
}
|
||||
|
||||
var _ store.PeerStore = &peerStore{}
|
||||
|
||||
func (s *peerStore) shardIndex(infohash string) uint32 {
|
||||
idx := fnv.New32()
|
||||
idx.Write([]byte(infohash))
|
||||
return idx.Sum32() % uint32(len(s.shards))
|
||||
}
|
||||
|
||||
func (s *peerStore) PutSeeder(infohash string, p chihaya.Peer) error {
|
||||
key := infohash + seedersSuffix
|
||||
|
||||
shard := s.shards[s.shardIndex(infohash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
||||
if shard.peers[key] == nil {
|
||||
shard.peers[key] = make(map[string]peer)
|
||||
}
|
||||
|
||||
shard.peers[key][p.ID] = peer{
|
||||
Peer: p,
|
||||
LastAction: time.Now(),
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *peerStore) DeleteSeeder(infohash, peerID string) error {
|
||||
key := infohash + seedersSuffix
|
||||
|
||||
shard := s.shards[s.shardIndex(infohash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
||||
if shard.peers[key] == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
delete(shard.peers[key], peerID)
|
||||
|
||||
if len(shard.peers[key]) == 0 {
|
||||
shard.peers[key] = nil
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *peerStore) PutLeecher(infohash string, p chihaya.Peer) error {
|
||||
key := infohash + leechersSuffix
|
||||
|
||||
shard := s.shards[s.shardIndex(infohash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
||||
if shard.peers[key] == nil {
|
||||
shard.peers[key] = make(map[string]peer)
|
||||
}
|
||||
|
||||
shard.peers[key][p.ID] = peer{
|
||||
Peer: p,
|
||||
LastAction: time.Now(),
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *peerStore) DeleteLeecher(infohash, peerID string) error {
|
||||
key := infohash + leechersSuffix
|
||||
|
||||
shard := s.shards[s.shardIndex(infohash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
||||
if shard.peers[key] == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
delete(shard.peers[key], peerID)
|
||||
|
||||
if len(shard.peers[key]) == 0 {
|
||||
shard.peers[key] = nil
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *peerStore) GraduateLeecher(infohash string, p chihaya.Peer) error {
|
||||
leecherKey := infohash + leechersSuffix
|
||||
seederKey := infohash + seedersSuffix
|
||||
|
||||
shard := s.shards[s.shardIndex(infohash)]
|
||||
shard.Lock()
|
||||
defer shard.Unlock()
|
||||
|
||||
if shard.peers[leecherKey] != nil {
|
||||
delete(shard.peers[leecherKey], p.ID)
|
||||
}
|
||||
|
||||
if shard.peers[seederKey] == nil {
|
||||
shard.peers[seederKey] = make(map[string]peer)
|
||||
}
|
||||
|
||||
shard.peers[seederKey][p.ID] = peer{
|
||||
Peer: p,
|
||||
LastAction: time.Now(),
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *peerStore) CollectGarbage(cutoff time.Time) error {
|
||||
for _, shard := range s.shards {
|
||||
shard.RLock()
|
||||
var keys []string
|
||||
for key := range shard.peers {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
shard.RUnlock()
|
||||
runtime.Gosched()
|
||||
|
||||
for _, key := range keys {
|
||||
shard.Lock()
|
||||
var peersToDelete []string
|
||||
for peerID, p := range shard.peers[key] {
|
||||
if p.LastAction.Before(cutoff) {
|
||||
peersToDelete = append(peersToDelete, peerID)
|
||||
}
|
||||
}
|
||||
|
||||
for _, peerID := range peersToDelete {
|
||||
delete(shard.peers[key], peerID)
|
||||
}
|
||||
shard.Unlock()
|
||||
runtime.Gosched()
|
||||
}
|
||||
|
||||
runtime.Gosched()
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *peerStore) AnnouncePeers(infohash string, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error) {
|
||||
leecherKey := infohash + leechersSuffix
|
||||
seederKey := infohash + seedersSuffix
|
||||
|
||||
shard := s.shards[s.shardIndex(infohash)]
|
||||
shard.RLock()
|
||||
defer shard.RUnlock()
|
||||
|
||||
if seeder {
|
||||
// Append leechers as possible.
|
||||
leechers := shard.peers[leecherKey]
|
||||
for _, p := range leechers {
|
||||
if numWant == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
if p.IP.To4() == nil {
|
||||
peers6 = append(peers6, p.Peer)
|
||||
} else {
|
||||
peers = append(peers, p.Peer)
|
||||
}
|
||||
numWant--
|
||||
}
|
||||
} else {
|
||||
// Append as many seeders as possible.
|
||||
seeders := shard.peers[seederKey]
|
||||
for _, p := range seeders {
|
||||
if numWant == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
if p.IP.To4() == nil {
|
||||
peers6 = append(peers6, p.Peer)
|
||||
} else {
|
||||
peers = append(peers, p.Peer)
|
||||
}
|
||||
numWant--
|
||||
}
|
||||
|
||||
// Append leechers until we reach numWant.
|
||||
leechers := shard.peers[leecherKey]
|
||||
if numWant > 0 {
|
||||
for _, p := range leechers {
|
||||
if numWant == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
if p.IP.To4() == nil {
|
||||
peers6 = append(peers6, p.Peer)
|
||||
} else {
|
||||
peers = append(peers, p.Peer)
|
||||
}
|
||||
numWant--
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package store
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/chihaya/chihaya"
|
||||
)
|
||||
|
||||
var peerStoreDrivers = make(map[string]PeerStoreDriver)
|
||||
|
||||
// PeerStore represents an interface for manipulating peers.
|
||||
type PeerStore interface {
|
||||
PutSeeder(infohash string, p chihaya.Peer) error
|
||||
DeleteSeeder(infohash, peerID string) error
|
||||
|
||||
PutLeecher(infohash string, p chihaya.Peer) error
|
||||
DeleteLeecher(infohash, peerID string) error
|
||||
|
||||
GraduateLeecher(infohash string, p chihaya.Peer) error
|
||||
AnnouncePeers(infohash string, seeder bool, numWant int) (peers, peers6 []chihaya.Peer, err error)
|
||||
CollectGarbage(cutoff time.Time) error
|
||||
}
|
||||
|
||||
// PeerStoreDriver represents an interface for creating a handle to the storage
|
||||
// of peers.
|
||||
type PeerStoreDriver interface {
|
||||
New(*Config) (PeerStore, error)
|
||||
}
|
||||
|
||||
// RegisterPeerStoreDriver makes a driver available by the provided name.
|
||||
//
|
||||
// If this function is called twice with the same name or if the driver is nil,
|
||||
// it panics.
|
||||
func RegisterPeerStoreDriver(name string, driver PeerStoreDriver) {
|
||||
if driver == nil {
|
||||
panic("storage: could not register nil PeerStoreDriver")
|
||||
}
|
||||
|
||||
if _, dup := peerStoreDrivers[name]; dup {
|
||||
panic("storage: could not register duplicate PeerStoreDriver: " + name)
|
||||
}
|
||||
|
||||
peerStoreDrivers[name] = driver
|
||||
}
|
||||
|
||||
// OpenPeerStore returns a PeerStore specified by a configuration.
|
||||
func OpenPeerStore(name string, cfg *Config) (PeerStore, error) {
|
||||
driver, ok := peerStoreDrivers[name]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf(
|
||||
"storage: unknown driver %q (forgotten import?)",
|
||||
name,
|
||||
)
|
||||
}
|
||||
|
||||
return driver.New(cfg)
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
// Copyright 2016 The Chihaya Authors. All rights reserved.
|
||||
// Use of this source code is governed by the BSD 2-Clause license,
|
||||
// which can be found in the LICENSE file.
|
||||
|
||||
package store
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"gopkg.in/yaml.v2"
|
||||
|
||||
"github.com/chihaya/chihaya/config"
|
||||
"github.com/chihaya/chihaya/server"
|
||||
"github.com/chihaya/chihaya/tracker"
|
||||
)
|
||||
|
||||
var theStore *Store
|
||||
|
||||
func init() {
|
||||
server.Register("store", constructor)
|
||||
}
|
||||
|
||||
func constructor(srvcfg *config.ServerConfig, tkr *tracker.Tracker) (server.Server, error) {
|
||||
if theStore == nil {
|
||||
cfg, err := newConfig(srvcfg)
|
||||
if err != nil {
|
||||
return nil, errors.New("store: invalid store config: " + err.Error())
|
||||
}
|
||||
|
||||
theStore = &Store{
|
||||
cfg: cfg,
|
||||
tkr: tkr,
|
||||
}
|
||||
}
|
||||
return theStore, nil
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
Addr string `yaml:"addr"`
|
||||
RequestTimeout time.Duration `yaml:"requestTimeout"`
|
||||
ReadTimeout time.Duration `yaml:"readTimeout"`
|
||||
WriteTimeout time.Duration `yaml:"writeTimeout"`
|
||||
GCAfter time.Duration `yaml:"gcAfter"`
|
||||
ClientStore string `yaml:"clientStore"`
|
||||
ClientStoreConfig interface{} `yaml:"clienStoreConfig"`
|
||||
PeerStore string `yaml:"peerStore"`
|
||||
PeerStoreConfig interface{} `yaml:"peerStoreConfig"`
|
||||
}
|
||||
|
||||
func newConfig(srvcfg interface{}) (*Config, error) {
|
||||
bytes, err := yaml.Marshal(srvcfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var cfg Config
|
||||
err = yaml.Unmarshal(bytes, &cfg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
// MustGetStore is used by middleware to access the store.
|
||||
//
|
||||
// This function calls log.Fatal if a server hasn't been already created by
|
||||
// the server package.
|
||||
func MustGetStore() *Store {
|
||||
if theStore == nil {
|
||||
log.Fatal("store middleware used without store server")
|
||||
}
|
||||
return theStore
|
||||
}
|
||||
|
||||
type Store struct {
|
||||
cfg *Config
|
||||
tkr *tracker.Tracker
|
||||
shutdown chan struct{}
|
||||
wg sync.WaitGroup
|
||||
|
||||
PeerStore
|
||||
ClientStore
|
||||
}
|
||||
|
||||
func (s *Store) Start() {
|
||||
}
|
||||
|
||||
func (s *Store) Stop() {
|
||||
close(s.shutdown)
|
||||
s.wg.Wait()
|
||||
}
|
||||
Reference in New Issue
Block a user