├── .gitignore ├── LICENSE ├── Makefile ├── README.md ├── bt ├── options.go └── worker.go ├── cmd └── dhtsearch │ ├── main.go │ └── tag.go ├── db ├── pgsql.go └── sqlite.go ├── dht ├── messages.go ├── node.go ├── options.go ├── packet.go ├── remote_node.go ├── routing_table.go ├── routing_table_test.go └── slab.go ├── go.mod ├── go.sum ├── krpc ├── krpc.go └── krpc_test.go └── models ├── infohash.go ├── infohash_test.go ├── peer.go ├── storage.go ├── tag.go └── torrent.go /.gitignore: -------------------------------------------------------------------------------- 1 | dhtsearch* 2 | config.toml 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2017 Felix Hanley 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | this software and associated documentation files (the "Software"), to deal in 7 | the Software without restriction, including without limitation the rights to 8 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 9 | of the Software, and to permit persons to whom the Software is furnished to do 10 | so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | 2 | VERSION ?= $(shell git describe --tags --always) 3 | SRC := $(shell find . -type f -name '*.go') 4 | FLAGS := --tags fts5 5 | PLAT := windows darwin linux freebsd openbsd 6 | BINARY := $(patsubst %,dist/%,$(shell find cmd/* -maxdepth 0 -type d -exec basename {} \;)) 7 | RELEASE := $(foreach os, $(PLAT), $(patsubst %,%-$(os), $(BINARY))) 8 | 9 | .PHONY: build 10 | build: sqlite $(BINARY) 11 | 12 | .PHONY: release 13 | release: $(RELEASE) 14 | 15 | dist/%: export GOOS=$(word 2,$(subst -, ,$*)) 16 | dist/%: bin=$(word 1,$(subst -, ,$*)) 17 | dist/%: $(SRC) $(shell find cmd/$(bin) -type f -name '*.go') 18 | go build -ldflags "-X main.version=$(VERSION)" $(FLAGS) \ 19 | -o $@ ./cmd/$(bin) 20 | 21 | sqlite: 22 | CGO_ENABLED=1 go get -u $(FLAGS) github.com/mattn/go-sqlite3 \ 23 | && go install $(FLAGS) github.com/mattn/go-sqlite3 24 | 25 | .PHONY: test 26 | test: 27 | go test -short -coverprofile=coverage.out ./... \ 28 | && go tool cover -func=coverage.out 29 | 30 | .PHONY: lint 31 | lint: ; go vet ./... 32 | 33 | .PHONY: clean 34 | clean: 35 | rm -f coverage* 36 | rm -rf dist 37 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DHT Search 2 | 3 | This is a Mainline DHT crawler and BitTorrent client which also provides an 4 | HTTP interface to query the indexed data. 5 | 6 | Distributed Hash Table (DHT) is a distributed system storing key/value pairs, 7 | in this case it is specifically Mainline DHT, the type used by BitTorrent 8 | clients. The crawler also implements a number of extensions which enable it to 9 | get the metadata for the torrent enabling indexing and later searching. 10 | 11 | The crawler joins the DHT network and listens to the conversations between 12 | nodes, keeping track of interesting packets. The most interesting packets are 13 | those where another node announces they have a torrent available. 14 | 15 | This BitTorrent client only downloads the torrent metadata. The actual files 16 | hosted by the remote nodes are not retrieved. 17 | 18 | ## Features 19 | 20 | - **Tagging** of torrents metadata is fetched. The torrent is tagged using a 21 | set of regular expressions matched against the torrent name and the files in 22 | the torrent. 23 | 24 | - **Filtering** can be done by tags. By default all torrents tagged 'adult' are 25 | not indexed. See the SkipTags option in the configuration file. 26 | 27 | - **Full Text Search** using PostgreSQL's or Sqlite's text search vectors. 28 | Torrent names are weighted more than file names. 29 | 30 | - **Statistics** for the crawler process are available when the HTTP server is 31 | enabled. Fetch the JSON from the `/status` endpoint. 32 | 33 | - **Custom tags** can be defined in the configuration file. 34 | 35 | - **Cross-platform** builds for Windows, Macos, Linux, FreeBSD, and OpenBSD 36 | 37 | ## Installation 38 | 39 | There is a Makefile for GNU make: 40 | 41 | ```shell 42 | $ make build 43 | ``` 44 | 45 | ## Usage 46 | 47 | You are going to need to sort out any port forwarding if you are behind NAT so 48 | remote nodes can get to yours. 49 | 50 | Configuration is done via a [TOML](https://github.com/toml-lang/toml) formatted 51 | file or via flags passed to the daemon. 52 | 53 | The following command line flags are available: 54 | 55 | -base-port int 56 | listen port (and first of multiple ports) (default 6881) 57 | -debug 58 | provide debug output 59 | -dsn string 60 | Database DSN (default "postgres://dht:dht@localhost/dht?sslmode=disable") 61 | -http-address string 62 | HTTP listen address:port (default "localhost:6880") 63 | -no-http 64 | no HTTP service 65 | -num-nodes int 66 | number of nodes to start (default 1) 67 | -quiet 68 | log only errors 69 | 70 | and the following "advanced" options: 71 | 72 | -max-bt-workers int 73 | max number of BT workers (default 256) 74 | -max-dht-workers int 75 | max number of DHT workers (default 256) 76 | -peer-cache-size int 77 | memory cache of seen peers (default 200) 78 | -routing-table-size int 79 | number of remote nodes in routing table (default 1000) 80 | -tcp-timeout int 81 | TCP timeout in seconds (default 10) 82 | -udp-timeout int 83 | UDP timeout in seconds (default 10) 84 | 85 | These options enable you to start a number of DHT nodes thus implementing a 86 | small scale [Sybil attack](https://en.wikipedia.org/wiki/Sybil_attack). The 87 | first DHT node will take the port specified and each subsequent port is for the 88 | following nodes. 89 | 90 | ## TODO 91 | 92 | - Enable rate limiting. 93 | - Improve our manners on the DHT network (replies etc.). 94 | - Improve the routing table implementation. 95 | - Add results pagination. 96 | - Add tests! 97 | -------------------------------------------------------------------------------- /bt/options.go: -------------------------------------------------------------------------------- 1 | package bt 2 | 3 | import ( 4 | "src.userspace.com.au/dhtsearch/models" 5 | "src.userspace.com.au/logger" 6 | ) 7 | 8 | type Option func(*Worker) error 9 | 10 | // SetOnNewTorrent sets the callback 11 | func SetOnNewTorrent(f func(models.Torrent)) Option { 12 | return func(w *Worker) error { 13 | w.OnNewTorrent = f 14 | return nil 15 | } 16 | } 17 | 18 | // SetOnBadPeer sets the callback 19 | func SetOnBadPeer(f func(models.Peer)) Option { 20 | return func(w *Worker) error { 21 | w.OnBadPeer = f 22 | return nil 23 | } 24 | } 25 | 26 | // SetPort sets the port to listen on 27 | func SetPort(p int) Option { 28 | return func(w *Worker) error { 29 | w.port = p 30 | return nil 31 | } 32 | } 33 | 34 | // SetIPv6 enables IPv6 35 | func SetIPv6(b bool) Option { 36 | return func(w *Worker) error { 37 | if b { 38 | w.family = "tcp6" 39 | } 40 | return nil 41 | } 42 | } 43 | 44 | // SetLogger sets the logger 45 | func SetLogger(l logger.Logger) Option { 46 | return func(w *Worker) error { 47 | w.log = l 48 | return nil 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /bt/worker.go: -------------------------------------------------------------------------------- 1 | package bt 2 | 3 | import ( 4 | "bytes" 5 | "encoding/binary" 6 | "errors" 7 | "fmt" 8 | "io" 9 | "io/ioutil" 10 | "net" 11 | "time" 12 | 13 | "src.userspace.com.au/dhtsearch/krpc" 14 | "src.userspace.com.au/dhtsearch/models" 15 | "src.userspace.com.au/go-bencode" 16 | "src.userspace.com.au/logger" 17 | ) 18 | 19 | const ( 20 | // MsgRequest marks a request message type 21 | MsgRequest = iota 22 | // MsgData marks a data message type 23 | MsgData 24 | // MsgReject marks a reject message type 25 | MsgReject 26 | // MsgExtended marks it as an extended message 27 | MsgExtended = 20 28 | ) 29 | 30 | const ( 31 | // BlockSize is 2 ^ 14 32 | BlockSize = 16384 33 | // MaxMetadataSize represents the max medata it can accept 34 | MaxMetadataSize = BlockSize * 1000 35 | // HandshakeBit represents handshake bit 36 | HandshakeBit = 0 37 | // TCPTimeout for BT connections 38 | TCPTimeout = 5 39 | ) 40 | 41 | var handshakePrefix = []byte{ 42 | 19, 66, 105, 116, 84, 111, 114, 114, 101, 110, 116, 32, 112, 114, 43 | 111, 116, 111, 99, 111, 108, 0, 0, 0, 0, 0, 16, 0, 1, 44 | } 45 | 46 | type Worker struct { 47 | pool chan chan models.Peer 48 | port int 49 | family string 50 | OnNewTorrent func(t models.Torrent) 51 | OnBadPeer func(p models.Peer) 52 | log logger.Logger 53 | } 54 | 55 | func NewWorker(pool chan chan models.Peer, opts ...Option) (*Worker, error) { 56 | var err error 57 | w := &Worker{ 58 | pool: pool, 59 | } 60 | 61 | // Set variadic options passed 62 | for _, option := range opts { 63 | err = option(w) 64 | if err != nil { 65 | return nil, err 66 | } 67 | } 68 | return w, nil 69 | } 70 | 71 | func (bt *Worker) Run() error { 72 | peerCh := make(chan models.Peer) 73 | 74 | for { 75 | // Signal we are ready for work 76 | bt.pool <- peerCh 77 | 78 | select { 79 | case p := <-peerCh: 80 | // Got work 81 | bt.log.Debug("worker got work", "peer", p) 82 | md, err := bt.fetchMetadata(p) 83 | if err != nil { 84 | bt.log.Debug("failed to fetch metadata", "error", err) 85 | if bt.OnBadPeer != nil { 86 | bt.OnBadPeer(p) 87 | } 88 | continue 89 | } 90 | t, err := models.TorrentFromMetadata(p.Infohash, md) 91 | if err != nil { 92 | bt.log.Warn("failed to load torrent", "error", err) 93 | continue 94 | } 95 | if bt.OnNewTorrent != nil { 96 | bt.OnNewTorrent(*t) 97 | } 98 | } 99 | } 100 | } 101 | 102 | // fetchMetadata fetchs medata info accroding to infohash from dht. 103 | func (bt *Worker) fetchMetadata(p models.Peer) (out []byte, err error) { 104 | var ( 105 | length int 106 | msgType byte 107 | totalPieces int 108 | pieces [][]byte 109 | utMetadata int 110 | metadataSize int 111 | ) 112 | 113 | defer func() { 114 | pieces = nil 115 | recover() 116 | }() 117 | 118 | //ll := bt.log.WithFields("address", p.Addr.String()) 119 | 120 | //ll.Debug("connecting") 121 | dial, err := net.DialTimeout("tcp", p.Addr.String(), time.Second*15) 122 | if err != nil { 123 | return out, err 124 | } 125 | // Cast 126 | conn := dial.(*net.TCPConn) 127 | conn.SetLinger(0) 128 | defer conn.Close() 129 | //ll.Debug("dialed") 130 | 131 | data := bytes.NewBuffer(nil) 132 | data.Grow(BlockSize) 133 | 134 | ih := models.GenInfohash() 135 | 136 | // TCP handshake 137 | //ll.Debug("sending handshake") 138 | _, err = sendHandshake(conn, p.Infohash, ih) 139 | if err != nil { 140 | return nil, err 141 | } 142 | 143 | // Handle the handshake response 144 | //ll.Debug("handling handshake response") 145 | err = read(conn, 68, data) 146 | if err != nil { 147 | return nil, err 148 | } 149 | next := data.Next(68) 150 | //ll.Debug("got next data") 151 | if !(bytes.Equal(handshakePrefix[:20], next[:20]) && next[25]&0x10 != 0) { 152 | //ll.Debug("next data does not match", "next", next) 153 | return nil, errors.New("invalid handshake response") 154 | } 155 | 156 | //ll.Debug("sending ext handshake") 157 | _, err = sendExtHandshake(conn) 158 | if err != nil { 159 | return nil, err 160 | } 161 | 162 | for { 163 | length, err = readMessage(conn, data) 164 | if err != nil { 165 | return out, err 166 | } 167 | 168 | if length == 0 { 169 | continue 170 | } 171 | 172 | msgType, err = data.ReadByte() 173 | if err != nil { 174 | return out, err 175 | } 176 | 177 | switch msgType { 178 | case MsgExtended: 179 | extendedID, err := data.ReadByte() 180 | if err != nil { 181 | return out, err 182 | } 183 | 184 | payload, err := ioutil.ReadAll(data) 185 | if err != nil { 186 | return out, err 187 | } 188 | 189 | if extendedID == 0 { 190 | if pieces != nil { 191 | return out, errors.New("invalid extended ID") 192 | } 193 | 194 | utMetadata, metadataSize, err = getUTMetaSize(payload) 195 | if err != nil { 196 | return out, err 197 | } 198 | 199 | totalPieces = metadataSize / BlockSize 200 | if metadataSize%BlockSize != 0 { 201 | totalPieces++ 202 | } 203 | 204 | pieces = make([][]byte, totalPieces) 205 | go bt.requestPieces(conn, utMetadata, metadataSize, totalPieces) 206 | 207 | continue 208 | } 209 | 210 | if pieces == nil { 211 | return out, errors.New("no pieces found") 212 | } 213 | 214 | dict, index, err := bencode.DecodeDict(payload, 0) 215 | if err != nil { 216 | return out, err 217 | } 218 | 219 | mt, err := krpc.GetInt(dict, "msg_type") 220 | if err != nil { 221 | return out, err 222 | } 223 | 224 | if mt != MsgData { 225 | continue 226 | } 227 | 228 | piece, err := krpc.GetInt(dict, "piece") 229 | if err != nil { 230 | return out, err 231 | } 232 | 233 | pieceLen := length - 2 - index 234 | 235 | // Not last piece? should be full block 236 | if totalPieces > 1 && piece != totalPieces-1 && pieceLen != BlockSize { 237 | return out, fmt.Errorf("incomplete piece %d", piece) 238 | } 239 | // Last piece needs to equal remainder 240 | if piece == totalPieces-1 && pieceLen != metadataSize%BlockSize { 241 | return out, fmt.Errorf("incorrect final piece %d", piece) 242 | } 243 | 244 | pieces[piece] = payload[index:] 245 | 246 | if bt.isDone(pieces) { 247 | return bytes.Join(pieces, nil), nil 248 | } 249 | default: 250 | data.Reset() 251 | } 252 | } 253 | } 254 | 255 | // isDone checks if all pieces are complete 256 | func (bt *Worker) isDone(pieces [][]byte) bool { 257 | for _, piece := range pieces { 258 | if len(piece) == 0 { 259 | return false 260 | } 261 | } 262 | return true 263 | } 264 | 265 | // read reads size-length bytes from conn to data. 266 | func read(conn net.Conn, size int, data io.Writer) error { 267 | conn.SetReadDeadline(time.Now().Add(time.Second * time.Duration(TCPTimeout))) 268 | n, err := io.CopyN(data, conn, int64(size)) 269 | if err != nil || n != int64(size) { 270 | return errors.New("read error") 271 | } 272 | return nil 273 | } 274 | 275 | // readMessage gets a message from the tcp connection. 276 | func readMessage(conn net.Conn, data *bytes.Buffer) (length int, err error) { 277 | if err = read(conn, 4, data); err != nil { 278 | return length, err 279 | } 280 | 281 | length, err = bytes2int(data.Next(4)) 282 | if err != nil { 283 | return length, err 284 | } 285 | 286 | if length == 0 { 287 | return length, nil 288 | } 289 | 290 | err = read(conn, length, data) 291 | return length, err 292 | } 293 | 294 | // sendMessage sends data to the connection. 295 | func sendMessage(conn net.Conn, data []byte) (int, error) { 296 | length := int32(len(data)) 297 | 298 | buffer := bytes.NewBuffer(nil) 299 | binary.Write(buffer, binary.BigEndian, length) 300 | 301 | conn.SetWriteDeadline(time.Now().Add(time.Second * time.Duration(TCPTimeout))) 302 | return conn.Write(append(buffer.Bytes(), data...)) 303 | } 304 | 305 | // sendHandshake sends handshake message to conn. 306 | func sendHandshake(conn net.Conn, ih, id models.Infohash) (int, error) { 307 | data := make([]byte, 68) 308 | copy(data[:28], handshakePrefix) 309 | copy(data[28:48], []byte(ih)) 310 | copy(data[48:], []byte(id)) 311 | 312 | conn.SetWriteDeadline(time.Now().Add(time.Second * time.Duration(TCPTimeout))) 313 | return conn.Write(data) 314 | } 315 | 316 | // onHandshake handles the handshake response. 317 | func onHandshake(data []byte) (err error) { 318 | if !(bytes.Equal(handshakePrefix[:20], data[:20]) && data[25]&0x10 != 0) { 319 | err = errors.New("invalid handshake response") 320 | } 321 | return err 322 | } 323 | 324 | // sendExtHandshake requests for the ut_metadata and metadata_size. 325 | func sendExtHandshake(conn net.Conn) (int, error) { 326 | m, err := bencode.EncodeDict(map[string]interface{}{ 327 | "m": map[string]interface{}{"ut_metadata": 1}, 328 | }) 329 | if err != nil { 330 | return 0, err 331 | } 332 | data := append([]byte{MsgExtended, HandshakeBit}, m...) 333 | 334 | return sendMessage(conn, data) 335 | } 336 | 337 | // getUTMetaSize returns the ut_metadata and metadata_size. 338 | func getUTMetaSize(data []byte) (utMetadata int, metadataSize int, err error) { 339 | dict, _, err := bencode.DecodeDict(data, 0) 340 | if err != nil { 341 | return utMetadata, metadataSize, err 342 | } 343 | 344 | m, err := krpc.GetMap(dict, "m") 345 | if err != nil { 346 | return utMetadata, metadataSize, err 347 | } 348 | 349 | utMetadata, err = krpc.GetInt(m, "ut_metadata") 350 | if err != nil { 351 | return utMetadata, metadataSize, err 352 | } 353 | 354 | metadataSize, err = krpc.GetInt(dict, "metadata_size") 355 | if err != nil { 356 | return utMetadata, metadataSize, err 357 | } 358 | 359 | if metadataSize > MaxMetadataSize { 360 | err = errors.New("metadata_size too long") 361 | } 362 | return utMetadata, metadataSize, err 363 | } 364 | 365 | // Request more pieces 366 | func (bt *Worker) requestPieces(conn net.Conn, utMetadata int, metadataSize int, totalPieces int) { 367 | buffer := make([]byte, 1024) 368 | for i := 0; i < totalPieces; i++ { 369 | buffer[0] = MsgExtended 370 | buffer[1] = byte(utMetadata) 371 | 372 | msg, _ := bencode.EncodeDict(map[string]interface{}{ 373 | "msg_type": MsgRequest, 374 | "piece": i, 375 | }) 376 | 377 | length := len(msg) + 2 378 | copy(buffer[2:length], msg) 379 | 380 | sendMessage(conn, buffer[:length]) 381 | } 382 | buffer = nil 383 | } 384 | 385 | // bytes2int returns the int value it represents. 386 | func bytes2int(data []byte) (int, error) { 387 | n := len(data) 388 | if n > 8 { 389 | return 0, errors.New("data too long") 390 | } 391 | 392 | val := uint64(0) 393 | 394 | for i, b := range data { 395 | val += uint64(b) << uint64((n-i-1)*8) 396 | } 397 | return int(val), nil 398 | } 399 | -------------------------------------------------------------------------------- /cmd/dhtsearch/main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "flag" 5 | "fmt" 6 | "os" 7 | "regexp" 8 | "strings" 9 | "time" 10 | "unicode" 11 | 12 | "github.com/hashicorp/golang-lru" 13 | "src.userspace.com.au/dhtsearch/bt" 14 | "src.userspace.com.au/dhtsearch/db" 15 | "src.userspace.com.au/dhtsearch/dht" 16 | "src.userspace.com.au/dhtsearch/models" 17 | "src.userspace.com.au/logger" 18 | //"github.com/pkg/profile" 19 | ) 20 | 21 | var ( 22 | version string 23 | log logger.Logger 24 | ) 25 | 26 | // DHT vars 27 | var ( 28 | debug bool 29 | port int 30 | ipv6 bool 31 | dhtNodes int 32 | showVersion bool 33 | ) 34 | 35 | // Torrent vars 36 | var ( 37 | pool chan chan models.Peer 38 | torrents chan models.Torrent 39 | btNodes int 40 | tagREs map[string]*regexp.Regexp 41 | skipTags string 42 | ) 43 | 44 | // Store vars 45 | var ( 46 | dsn string 47 | ihBlacklist *lru.ARCCache 48 | peerBlacklist *lru.ARCCache 49 | ) 50 | 51 | func main() { 52 | //defer profile.Start(profile.MemProfile).Stop() 53 | flag.IntVar(&port, "port", 6881, "listen port (and first for multiple nodes") 54 | flag.BoolVar(&debug, "debug", false, "show debug output") 55 | flag.BoolVar(&ipv6, "6", false, "listen on IPv6 also") 56 | flag.IntVar(&dhtNodes, "dht-nodes", 1, "number of DHT nodes to start") 57 | 58 | flag.IntVar(&btNodes, "bt-nodes", 3, "number of BT nodes to start") 59 | flag.StringVar(&skipTags, "skip-tags", "xxx", "tags of torrents to skip") 60 | 61 | flag.StringVar(&dsn, "dsn", "file:dhtsearch.db?cache=shared&mode=memory", "database DSN") 62 | 63 | flag.BoolVar(&showVersion, "v", false, "show version") 64 | 65 | flag.Parse() 66 | 67 | if showVersion { 68 | fmt.Println(version) 69 | os.Exit(0) 70 | } 71 | 72 | logOpts := &logger.Options{ 73 | Name: "dhtsearch", 74 | Level: logger.Info, 75 | } 76 | 77 | if debug { 78 | logOpts.Level = logger.Debug 79 | } 80 | log = logger.New(logOpts) 81 | log.Info("version", version) 82 | log.Debug("debugging") 83 | 84 | store, err := db.NewStore(dsn) 85 | if err != nil { 86 | log.Error("failed to connect store", "error", err) 87 | os.Exit(1) 88 | } 89 | defer store.Close() 90 | 91 | createTagRegexps() 92 | 93 | ihBlacklist, err = lru.NewARC(1000) 94 | if err != nil { 95 | log.Error("failed to create infohash blacklist", "error", err) 96 | os.Exit(1) 97 | } 98 | peerBlacklist, err = lru.NewARC(1000) 99 | if err != nil { 100 | log.Error("failed to create blacklist", "error", err) 101 | os.Exit(1) 102 | } 103 | // TODO read in existing blacklist 104 | // TODO populate bloom filter 105 | 106 | go startDHTNodes(store) 107 | 108 | go startBTWorkers(store) 109 | 110 | go processPendingPeers(store) 111 | 112 | for { 113 | select { 114 | case <-time.After(300 * time.Second): 115 | log.Info("---- mark ----") 116 | } 117 | } 118 | } 119 | 120 | func startDHTNodes(s models.PeerStore) { 121 | log.Debug("starting dht nodes") 122 | nodes := make([]*dht.Node, dhtNodes) 123 | 124 | for i := 0; i < dhtNodes; i++ { 125 | dht, err := dht.NewNode( 126 | dht.SetLogger(log.Named("dht")), 127 | dht.SetPort(port+i), 128 | dht.SetIPv6(ipv6), 129 | dht.SetBlacklist(peerBlacklist), 130 | dht.SetOnAnnouncePeer(func(p models.Peer) { 131 | if _, black := ihBlacklist.Get(p.Infohash.String()); black { 132 | log.Debug("ignoring blacklisted infohash", "peer", p) 133 | return 134 | } 135 | //log.Debug("peer announce", "peer", p) 136 | err := s.SavePeer(&p) 137 | if err != nil { 138 | log.Error("failed to save peer", "error", err) 139 | } 140 | }), 141 | dht.SetOnBadPeer(func(p models.Peer) { 142 | err := s.RemovePeer(&p) 143 | if err != nil { 144 | log.Error("failed to remove peer", "error", err) 145 | } 146 | }), 147 | ) 148 | if err != nil { 149 | log.Error("failed to create node", "error", err) 150 | continue 151 | } 152 | go dht.Run() 153 | nodes[i] = dht 154 | } 155 | } 156 | 157 | func processPendingPeers(s models.InfohashStore) { 158 | log.Debug("processing pending peers") 159 | for { 160 | peers, err := s.PendingInfohashes(10) 161 | if err != nil { 162 | log.Warn("failed to get pending peer", "error", err) 163 | time.Sleep(time.Second * 1) 164 | continue 165 | } 166 | for _, p := range peers { 167 | log.Debug("pending peer retrieved", "peer", *p) 168 | select { 169 | case w := <-pool: 170 | //log.Debug("assigning peer to bt worker") 171 | w <- *p 172 | } 173 | } 174 | } 175 | } 176 | 177 | func startBTWorkers(s models.TorrentStore) { 178 | log.Debug("starting bittorrent workers") 179 | pool = make(chan chan models.Peer) 180 | torrents = make(chan models.Torrent) 181 | 182 | onNewTorrent := func(t models.Torrent) { 183 | // Add tags 184 | tags := tagTorrent(t, tagREs) 185 | for _, skipTag := range strings.Split(skipTags, ",") { 186 | for _, tg := range tags { 187 | if skipTag == tg { 188 | log.Debug("skipping torrent", "infohash", t.Infohash, "tags", tags) 189 | ihBlacklist.Add(t.Infohash.String(), true) 190 | s.RemoveTorrent(&t) 191 | return 192 | } 193 | } 194 | } 195 | t.Tags = tags 196 | log.Debug("torrent tagged", "infohash", t.Infohash, "tags", tags) 197 | err := s.SaveTorrent(&t) 198 | if err != nil { 199 | log.Error("failed to save torrent", "error", err) 200 | ihBlacklist.Add(t.Infohash.String(), true) 201 | s.RemoveTorrent(&t) 202 | } 203 | log.Info("torrent added", "name", t.Name, "size", t.Size, "tags", t.Tags) 204 | } 205 | 206 | onBadPeer := func(p models.Peer) { 207 | log.Debug("removing peer", "peer", p) 208 | err := s.RemovePeer(&p) 209 | if err != nil { 210 | log.Error("failed to remove peer", "peer", p, "error", err) 211 | } 212 | peerBlacklist.Add(p.Addr.String(), true) 213 | } 214 | 215 | for i := 0; i < btNodes; i++ { 216 | w, err := bt.NewWorker( 217 | pool, 218 | bt.SetLogger(log.Named("bt")), 219 | bt.SetPort(port+i), 220 | bt.SetIPv6(ipv6), 221 | bt.SetOnNewTorrent(onNewTorrent), 222 | bt.SetOnBadPeer(onBadPeer), 223 | ) 224 | if err != nil { 225 | log.Error("failed to create bt worker", "error", err) 226 | return 227 | } 228 | log.Debug("running bt node", "index", i) 229 | go w.Run() 230 | } 231 | } 232 | 233 | // Filter on words, existing 234 | func createTagRegexps() { 235 | tagREs = make(map[string]*regexp.Regexp) 236 | for tag, re := range tags { 237 | tagREs[tag] = regexp.MustCompile("(?i)" + re) 238 | } 239 | // Add character classes 240 | for cc, _ := range unicode.Scripts { 241 | if cc == "Latin" || cc == "Common" { 242 | continue 243 | } 244 | className := strings.ToLower(cc) 245 | // Test for 3 or more characters per character class 246 | tagREs[className] = regexp.MustCompile(fmt.Sprintf(`(?i)\p{%s}{3,}`, cc)) 247 | } 248 | // Merge user tags 249 | /* 250 | for tag, re := range Config.Tags { 251 | if !Config.Quiet { 252 | fmt.Printf("Adding user tag: %s = %s\n", tag, re) 253 | } 254 | tagREs[tag] = regexp.MustCompile("(?i)" + re) 255 | } 256 | */ 257 | } 258 | -------------------------------------------------------------------------------- /cmd/dhtsearch/tag.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "fmt" 5 | "regexp" 6 | "strings" 7 | "unicode" 8 | 9 | "src.userspace.com.au/dhtsearch/models" 10 | ) 11 | 12 | // Default tags, can be supplimented or overwritten by config 13 | var tags = map[string]string{ 14 | "flac": `\.flac$`, 15 | "episode": "(season|episode|s[0-9]{2}e[0-9]{2})", 16 | "1080": "1080", 17 | "720": "720", 18 | "hd": "hd|720|1080|4k", 19 | "rip": "(bdrip|dvd[- ]?rip|dvdmux|br[- ]?rip|dvd[-]?r|web[- ]?dl|hdrip)", 20 | "xxx": `(xxx|p(orn|ussy)|censor|sex|urbat|a(ss|nal)|o(rgy|gasm)|(fu|di|co)ck|esbian|milf|lust|gay)|rotic|18(\+|yr)|hore|hemale|virgin`, 21 | "ebook": "epub", 22 | "application": `\.(apk|exe|msi|dmg)$`, 23 | "android": `\.apk$`, 24 | "apple": `\.dmg$`, 25 | "subtitles": `\.s(rt|ub)$`, 26 | "archive": `\.(zip|rar|p7|tgz|bz2|iso)$`, 27 | "video": `(\.(3g2|3gp|amv|asf|avi|drc|f4a|f4b|f4p|f4v|flv|gif|gifv|m2v|m4p|m4v|mkv|mng|mov|mp2|mp4|mpe|mpeg|mpg|mpv|mxf|net|nsv|ogv|qt|rm|rmvb|roq|svi|vob|webm|wmv|yuv)$|divx|x264|x265)`, 28 | "audio": `\.(aa|aac|aax|act|aiff|amr|ape|au|awb|dct|dss|dvf|flac|gsm|iklax|ivs|m4a|m4b|mmf|mp3|mpc|msv|ogg|opus|ra|raw|sln|tta|vox|wav|wma|wv)$`, 29 | "document": `\.(cbr|cbz|cb7|cbt|cba|epub|djvu|fb2|ibook|azw.|lit|prc|mobi|pdb|pdb|oxps|xps|pdf)$`, 30 | "bootleg": `(camrip|hdts|[-. ](ts|tc)[-. ]|hdtc)`, 31 | "screener": `(bd[-]?scr|screener|dvd[-]?scr|r5)`, 32 | "font": `(font|\.(ttf|fon|otf)$)`, 33 | } 34 | 35 | func mergeCharacterTagREs(tagREs map[string]*regexp.Regexp) error { 36 | var err error 37 | // Add character classes 38 | for cc := range unicode.Scripts { 39 | if cc == "Latin" || cc == "Common" { 40 | continue 41 | } 42 | className := strings.ToLower(cc) 43 | // Test for 3 or more characters per character class 44 | tagREs[className], err = regexp.Compile(fmt.Sprintf(`(?i)\p{%s}{3,}`, cc)) 45 | if err != nil { 46 | return err 47 | } 48 | } 49 | return nil 50 | } 51 | 52 | func mergeTagRegexps(tagREs map[string]*regexp.Regexp, tags map[string]string) error { 53 | var err error 54 | for tag, re := range tags { 55 | tagREs[tag], err = regexp.Compile("(?i)" + re) 56 | if err != nil { 57 | return err 58 | } 59 | } 60 | return nil 61 | } 62 | 63 | func tagTorrent(t models.Torrent, tagREs map[string]*regexp.Regexp) (tags []string) { 64 | ttags := make(map[string]bool) 65 | 66 | for tag, re := range tagREs { 67 | if re.MatchString(t.Name) { 68 | ttags[tag] = true 69 | } 70 | for _, f := range t.Files { 71 | if re.MatchString(f.Path) { 72 | ttags[tag] = true 73 | } 74 | } 75 | } 76 | // Make unique 77 | for tt := range ttags { 78 | tags = append(tags, tt) 79 | } 80 | return tags 81 | } 82 | 83 | func hasTag(t models.Torrent, tag string) bool { 84 | for _, t := range t.Tags { 85 | if tag == t { 86 | return true 87 | } 88 | } 89 | return false 90 | } 91 | -------------------------------------------------------------------------------- /db/pgsql.go: -------------------------------------------------------------------------------- 1 | // +build ignore postgres 2 | 3 | package db 4 | 5 | import ( 6 | "fmt" 7 | "net" 8 | 9 | "github.com/jackc/pgx" 10 | "github.com/jackc/pgx/pgtype" 11 | "src.userspace.com.au/dhtsearch/models" 12 | ) 13 | 14 | // Store is a store 15 | type Store struct { 16 | *pgx.ConnPool 17 | } 18 | 19 | // NewStore connects and initializes a new store 20 | func NewStore(dsn string) (*Store, error) { 21 | cfg, err := pgx.ParseURI(dsn) 22 | if err != nil { 23 | return nil, err 24 | } 25 | c, err := pgx.NewConnPool(pgx.ConnPoolConfig{ConnConfig: cfg, MaxConnections: 10}) 26 | if err != nil { 27 | return nil, err 28 | } 29 | 30 | s := &Store{c} 31 | 32 | err = s.migrate() 33 | if err != nil { 34 | return nil, err 35 | } 36 | 37 | err = s.prepareStatements() 38 | if err != nil { 39 | return nil, err 40 | } 41 | 42 | return s, err 43 | } 44 | 45 | // PendingInfohashes gets the next pending infohash from the store 46 | func (s *Store) PendingInfohashes(n int) (peers []*models.Peer, err error) { 47 | 48 | rows, err := s.Query("selectPendingInfohashes", n) 49 | if err != nil { 50 | return nil, err 51 | } 52 | defer rows.Close() 53 | for rows.Next() { 54 | var p models.Peer 55 | var ih pgtype.Bytea 56 | var addr string 57 | err = rows.Scan(&addr, &ih) 58 | if err != nil { 59 | return nil, err 60 | } 61 | // TODO save peer network? 62 | p.Addr, err = net.ResolveUDPAddr("udp", addr) 63 | if err != nil { 64 | return nil, err 65 | } 66 | ih.AssignTo(&p.Infohash) 67 | peers = append(peers, &p) 68 | } 69 | return peers, nil 70 | } 71 | 72 | // SaveTorrent implements torrentStore 73 | func (s *Store) SaveTorrent(t *models.Torrent) error { 74 | tx, err := s.Begin() 75 | if err != nil { 76 | return err 77 | } 78 | defer tx.Rollback() 79 | 80 | var torrentID int 81 | err = tx.QueryRow("insertTorrent", t.Name, t.Infohash, t.Size).Scan(&torrentID) 82 | if err != nil { 83 | return fmt.Errorf("insertTorrent: %s", err) 84 | } 85 | 86 | // Write tags 87 | for _, tag := range t.Tags { 88 | tagID, err := s.SaveTag(tag) 89 | if err != nil { 90 | return fmt.Errorf("saveTag: %s", err) 91 | } 92 | _, err = tx.Exec("insertTagTorrent", tagID, torrentID) 93 | if err != nil { 94 | return fmt.Errorf("insertTagTorrent: %s", err) 95 | } 96 | } 97 | 98 | // Write files 99 | for _, f := range t.Files { 100 | _, err := tx.Exec("insertFile", torrentID, f.Path, f.Size) 101 | if err != nil { 102 | return fmt.Errorf("insertFile: %s", err) 103 | } 104 | } 105 | 106 | // Should this be outside the transaction? 107 | _, err = tx.Exec("updateFTSVectors", torrentID) 108 | if err != nil { 109 | return fmt.Errorf("updateVectors: %s", err) 110 | } 111 | return tx.Commit() 112 | } 113 | 114 | func (s *Store) RemoveTorrent(t *models.Torrent) (err error) { 115 | _, err = s.Exec("removeTorrent", t.Infohash.Bytes()) 116 | return err 117 | } 118 | 119 | // SavePeer implements torrentStore 120 | func (s *Store) SavePeer(p *models.Peer) (err error) { 121 | _, err = s.Exec("insertPeer", p.Addr.String(), p.Infohash.Bytes()) 122 | return err 123 | } 124 | 125 | func (s *Store) RemovePeer(p *models.Peer) (err error) { 126 | _, err = s.Exec("removePeer", p.Addr.String()) 127 | return err 128 | } 129 | 130 | // TorrentsByHash implements torrentStore 131 | func (s *Store) TorrentByHash(ih models.Infohash) (*models.Torrent, error) { 132 | rows, err := s.Query("getTorrent", ih) 133 | if err != nil { 134 | return nil, err 135 | } 136 | defer rows.Close() 137 | torrents, err := s.fetchTorrents(rows) 138 | if err != nil { 139 | return nil, err 140 | } 141 | return torrents[0], nil 142 | } 143 | 144 | // TorrentsByName implements torrentStore 145 | func (s *Store) TorrentsByName(query string, offset int) ([]*models.Torrent, error) { 146 | rows, err := s.Query("searchTorrents", fmt.Sprintf("%%%s%%", query), offset) 147 | if err != nil { 148 | return nil, err 149 | } 150 | defer rows.Close() 151 | torrents, err := s.fetchTorrents(rows) 152 | if err != nil { 153 | return nil, err 154 | } 155 | return torrents, nil 156 | } 157 | 158 | // TorrentsByTag implements torrentStore 159 | func (s *Store) TorrentsByTag(tag string, offset int) ([]*models.Torrent, error) { 160 | rows, err := s.Query("torrentsByTag", tag, offset) 161 | if err != nil { 162 | return nil, err 163 | } 164 | defer rows.Close() 165 | torrents, err := s.fetchTorrents(rows) 166 | if err != nil { 167 | return nil, err 168 | } 169 | return torrents, nil 170 | } 171 | 172 | // SaveTag implements tagStore interface 173 | func (s *Store) SaveTag(tag string) (tagID int, err error) { 174 | err = s.QueryRow("insertTag", tag).Scan(&tagID) 175 | return tagID, err 176 | } 177 | 178 | func (s *Store) fetchTorrents(rows *pgx.Rows) (torrents []*models.Torrent, err error) { 179 | for rows.Next() { 180 | var t models.Torrent 181 | /* 182 | t := &models.Torrent{ 183 | Files: []models.File{}, 184 | Tags: []string{}, 185 | } 186 | */ 187 | err = rows.Scan( 188 | &t.ID, &t.Infohash, &t.Name, &t.Size, &t.Created, &t.Updated, 189 | ) 190 | if err != nil { 191 | return nil, err 192 | } 193 | 194 | err = func() error { 195 | rowsf, err := s.Query("selectFiles", t.ID) 196 | defer rowsf.Close() 197 | if err != nil { 198 | return fmt.Errorf("failed to select files: %s", err) 199 | } 200 | for rowsf.Next() { 201 | var f models.File 202 | err = rowsf.Scan(&f.ID, &f.TorrentID, &f.Path, &f.Size) 203 | if err != nil { 204 | return fmt.Errorf("failed to build file: %s", err) 205 | } 206 | } 207 | return nil 208 | }() 209 | if err != nil { 210 | return nil, err 211 | } 212 | 213 | err = func() error { 214 | rowst, err := s.Query("selectTags", t.ID) 215 | defer rowst.Close() 216 | if err != nil { 217 | return fmt.Errorf("failed to select tags: %s", err) 218 | } 219 | for rowst.Next() { 220 | var tg string 221 | err = rowst.Scan(&tg) 222 | if err != nil { 223 | return fmt.Errorf("failed to build tag: %s", err) 224 | } 225 | t.Tags = append(t.Tags, tg) 226 | } 227 | return nil 228 | }() 229 | if err != nil { 230 | return nil, err 231 | } 232 | torrents = append(torrents, &t) 233 | } 234 | return torrents, err 235 | } 236 | 237 | func (s *Store) migrate() error { 238 | tx, err := s.Begin() 239 | if err != nil { 240 | return err 241 | } 242 | defer tx.Rollback() 243 | 244 | var initialized bool 245 | err = tx.QueryRow(`select exists ( 246 | select 1 from pg_tables 247 | where schemaname = 'public' 248 | and tablename = 'settings' 249 | )`).Scan(&initialized) 250 | if err != nil { 251 | return err 252 | } 253 | 254 | if !initialized { 255 | _, err = tx.Exec("baseSchema") 256 | } 257 | 258 | // Start migrations 259 | var currentVersion int 260 | err = tx.QueryRow("select schema_version from settings").Scan(currentVersion) 261 | if err != nil { 262 | return err 263 | } 264 | 265 | switch currentVersion { 266 | case 1: 267 | default: 268 | } 269 | 270 | return nil 271 | } 272 | 273 | func (s *Store) prepareStatements() error { 274 | if _, err := s.Prepare( 275 | "removeTorrent", 276 | `delete from torrents 277 | where infohash = $1`, 278 | ); err != nil { 279 | return err 280 | } 281 | if _, err := s.Prepare( 282 | "selectPendingInfohashes", 283 | `with get_order as ( 284 | select t.id as torrent_id, min(pt.peer_id) as peer_id, count(pt.peer_id) as c 285 | from torrents t 286 | join peers_torrents pt on pt.torrent_id = t.id 287 | where t.name is null 288 | group by t.id 289 | -- order by c desc 290 | order by t.updated desc 291 | limit $1 292 | ) select p.address, t.infohash 293 | from get_order go 294 | join torrents t on t.id = go.torrent_id 295 | join peers p on p.id = go.peer_id`, 296 | ); err != nil { 297 | return err 298 | } 299 | 300 | if _, err := s.Prepare( 301 | "selectFiles", 302 | `select * from files 303 | where torrent_id = $1 304 | order by path asc`, 305 | ); err != nil { 306 | return err 307 | } 308 | 309 | if _, err := s.Prepare( 310 | "insertPeer", 311 | `with save_peer as ( 312 | insert into peers 313 | (address, created, updated) values ($1, now(), now()) 314 | returning id 315 | ), save_torrent as ( 316 | insert into torrents (infohash, created, updated) 317 | values ($2, now(), now()) 318 | on conflict (infohash) do update set 319 | updated = now() 320 | returning id 321 | ) insert into peers_torrents 322 | (peer_id, torrent_id) 323 | select 324 | sp.id, st.id 325 | from save_peer sp, save_torrent st 326 | on conflict do nothing`, 327 | ); err != nil { 328 | return err 329 | } 330 | 331 | if _, err := s.Prepare( 332 | "getTorrent", 333 | `select * from torrents where infohash = $1 limit 1`, 334 | ); err != nil { 335 | return err 336 | } 337 | 338 | if _, err := s.Prepare( 339 | "insertFile", 340 | `insert into files 341 | (torrent_id, path, size) 342 | values 343 | ($1, $2, $3)`, 344 | ); err != nil { 345 | return err 346 | } 347 | 348 | if _, err := s.Prepare( 349 | "selectTags", 350 | `select name 351 | from tags t 352 | inner join tags_torrents tt on t.id = tt.tag_id 353 | where tt.torrent_id = $1`, 354 | ); err != nil { 355 | return err 356 | } 357 | 358 | if _, err := s.Prepare( 359 | "removePeer", 360 | `delete from peers where address = $1`, 361 | ); err != nil { 362 | return err 363 | } 364 | 365 | if _, err := s.Prepare( 366 | "insertTagTorrent", 367 | `insert into tags_torrents 368 | (tag_id, torrent_id) values ($1, $2) 369 | on conflict do nothing`, 370 | ); err != nil { 371 | return err 372 | } 373 | 374 | if _, err := s.Prepare( 375 | "insertTag", 376 | `insert into tags (name) values ($1) 377 | on conflict (name) do update set name = excluded.name returning id`, 378 | ); err != nil { 379 | return err 380 | } 381 | 382 | if _, err := s.Prepare( 383 | "torrentsByTag", 384 | `select t.id, t.infohash, t.name, t.size, t.created, t.updated 385 | from torrents t 386 | inner join tags_torrents tt on t.id = tt.torrent_id 387 | inner join tags ta on tt.tag_id = ta.id 388 | where ta.name = $1 group by t.id 389 | order by updated asc 390 | limit 50 offset $2`, 391 | ); err != nil { 392 | return err 393 | } 394 | 395 | if _, err := s.Prepare( 396 | "searchTorrents", 397 | `select t.id, t.infohash, t.name, t.size, t.updated 398 | from torrents t 399 | where t.tsv @@ plainto_tsquery($1) 400 | order by ts_rank(tsv, plainto_tsquery($1)) desc, t.updated desc 401 | limit 50 offset $2`, 402 | ); err != nil { 403 | return err 404 | } 405 | 406 | if _, err := s.Prepare( 407 | "insertTorrent", 408 | `insert into torrents ( 409 | name, infohash, size, created, updated 410 | ) values ( 411 | $1, $2, $3, now(), now() 412 | ) on conflict (infohash) do 413 | update set 414 | name = $1, 415 | size = $3, 416 | updated = now() 417 | returning id`, 418 | ); err != nil { 419 | return err 420 | } 421 | 422 | if _, err := s.Prepare( 423 | "updateFTSVectors", 424 | `update torrents set 425 | tsv = sub.tsv from ( 426 | select t.id, 427 | setweight(to_tsvector( 428 | translate(t.name, '._-', ' ') 429 | ), 'A') 430 | || setweight(to_tsvector( 431 | translate(string_agg(coalesce(f.path, ''), ' '), './_-', ' ') 432 | ), 'B') as tsv 433 | from torrents t 434 | left join files f on t.id = f.torrent_id 435 | where t.id = $1 436 | group by t.id 437 | ) as sub 438 | where sub.id = torrents.id`, 439 | ); err != nil { 440 | return err 441 | } 442 | 443 | if _, err := s.Prepare( 444 | "baseSchema", 445 | `create table if not exists torrents ( 446 | id serial primary key, 447 | infohash bytea not null unique, 448 | size bigint, 449 | name text, 450 | created timestamp with time zone, 451 | updated timestamp with time zone, 452 | tsv tsvector 453 | ); 454 | create index tsv_idx on torrents using gin(tsv); 455 | create table if not exists files ( 456 | id serial not null primary key, 457 | torrent_id integer not null references torrents on delete cascade, 458 | path text, 459 | size bigint 460 | ); 461 | create table if not exists tags ( 462 | id serial primary key, 463 | name character varying(50) unique 464 | ); 465 | create table if not exists tags_torrents ( 466 | tag_id integer not null references tags (id) on delete cascade, 467 | torrent_id integer not null references torrents (id) on delete cascade, 468 | primary key (tag_id, torrent_id) 469 | ); 470 | create table if not exists peers ( 471 | id serial primary key, 472 | address character varying(50) not null, 473 | created timestamp with time zone, 474 | updated timestamp with time zone 475 | ); 476 | create table if not exists peers_torrents ( 477 | peer_id integer not null references peers (id) on delete cascade, 478 | torrent_id integer not null references torrents (id) on delete cascade, 479 | primary key (peer_id, torrent_id) 480 | ); 481 | create table if not exists settings ( 482 | schema_version integer not null 483 | ); 484 | insert into settings (schema_version) values (1);`, 485 | ); err != nil { 486 | return err 487 | } 488 | return nil 489 | } 490 | -------------------------------------------------------------------------------- /db/sqlite.go: -------------------------------------------------------------------------------- 1 | package db 2 | 3 | import ( 4 | "database/sql" 5 | "fmt" 6 | "net" 7 | "sync" 8 | 9 | _ "github.com/mattn/go-sqlite3" 10 | "src.userspace.com.au/dhtsearch/models" 11 | ) 12 | 13 | // Store is a store 14 | type Store struct { 15 | stmts map[string]*sql.Stmt 16 | conn *sql.DB 17 | lock sync.RWMutex 18 | } 19 | 20 | // NewStore connects and initializes a new store 21 | func NewStore(dsn string) (*Store, error) { 22 | conn, err := sql.Open("sqlite3", dsn) 23 | if err != nil { 24 | return nil, fmt.Errorf("failed to open store: %s", err) 25 | } 26 | 27 | s := &Store{conn: conn, stmts: make(map[string]*sql.Stmt)} 28 | 29 | err = s.migrate() 30 | if err != nil { 31 | return nil, err 32 | } 33 | 34 | err = s.prepareStatements() 35 | if err != nil { 36 | return nil, err 37 | } 38 | 39 | return s, err 40 | } 41 | 42 | func (s *Store) Close() error { 43 | return s.conn.Close() 44 | } 45 | 46 | // PendingInfohashes gets the next pending infohash from the store 47 | func (s *Store) PendingInfohashes(n int) (peers []*models.Peer, err error) { 48 | s.lock.RLock() 49 | defer s.lock.RUnlock() 50 | 51 | rows, err := s.stmts["selectPendingInfohashes"].Query(n) 52 | if err != nil { 53 | return nil, err 54 | } 55 | defer rows.Close() 56 | for rows.Next() { 57 | var p models.Peer 58 | var ih models.Infohash 59 | var addr string 60 | err = rows.Scan(&addr, &ih) 61 | if err != nil { 62 | return nil, err 63 | } 64 | // TODO save peer network? 65 | p.Addr, err = net.ResolveUDPAddr("udp", addr) 66 | if err != nil { 67 | return nil, err 68 | } 69 | p.Infohash = ih 70 | peers = append(peers, &p) 71 | } 72 | return peers, nil 73 | } 74 | 75 | // SaveTorrent implements torrentStore 76 | func (s *Store) SaveTorrent(t *models.Torrent) error { 77 | s.lock.Lock() 78 | defer s.lock.Unlock() 79 | 80 | tx, err := s.conn.Begin() 81 | if err != nil { 82 | return fmt.Errorf("saveTorrent: %s", err) 83 | } 84 | defer tx.Rollback() 85 | 86 | var torrentID int64 87 | var res sql.Result 88 | res, err = tx.Stmt(s.stmts["insertTorrent"]).Exec(t.Name, t.Infohash.Bytes(), t.Size) 89 | if err != nil { 90 | return fmt.Errorf("insertTorrent: %s", err) 91 | } 92 | if torrentID, err = res.LastInsertId(); err != nil { 93 | return fmt.Errorf("insertTorrent: %s", err) 94 | } 95 | 96 | // Write tags 97 | for _, tag := range t.Tags { 98 | var tagID int64 99 | 100 | res, err = tx.Stmt(s.stmts["insertTag"]).Exec(tag) 101 | if err != nil { 102 | return fmt.Errorf("saveTag: %s", err) 103 | } 104 | tagID, err = res.LastInsertId() 105 | if err != nil { 106 | return fmt.Errorf("saveTag: %s", err) 107 | } 108 | _, err = tx.Stmt(s.stmts["insertTagTorrent"]).Exec(tagID, torrentID) 109 | if err != nil { 110 | return fmt.Errorf("insertTagTorrent: %s", err) 111 | } 112 | } 113 | 114 | // Write files 115 | for _, f := range t.Files { 116 | _, err := tx.Stmt(s.stmts["insertFile"]).Exec(torrentID, f.Path, f.Size) 117 | if err != nil { 118 | return fmt.Errorf("insertFile: %s", err) 119 | } 120 | } 121 | 122 | return tx.Commit() 123 | } 124 | 125 | func (s *Store) RemoveTorrent(t *models.Torrent) (err error) { 126 | s.lock.Lock() 127 | defer s.lock.Unlock() 128 | 129 | _, err = s.stmts["removeTorrent"].Exec(t.Infohash) 130 | return fmt.Errorf("removeTorrent: %s", err) 131 | } 132 | 133 | // SavePeer implements torrentStore 134 | func (s *Store) SavePeer(p *models.Peer) (err error) { 135 | s.lock.Lock() 136 | defer s.lock.Unlock() 137 | 138 | var peerID int64 139 | var torrentID int64 140 | var res sql.Result 141 | 142 | tx, err := s.conn.Begin() 143 | if err != nil { 144 | return err 145 | } 146 | defer tx.Rollback() 147 | 148 | if res, err = tx.Stmt(s.stmts["insertPeer"]).Exec(p.Addr.String()); err != nil { 149 | return fmt.Errorf("savePeer: %s", err) 150 | } 151 | if peerID, err = res.LastInsertId(); err != nil { 152 | return fmt.Errorf("savePeer: %s", err) 153 | } 154 | 155 | if res, err = tx.Stmt(s.stmts["insertTorrent"]).Exec(nil, p.Infohash, 0); err != nil { 156 | return fmt.Errorf("savePeer: %s", err) 157 | } 158 | if torrentID, err = res.LastInsertId(); err != nil { 159 | return fmt.Errorf("savePeer: %s", err) 160 | } 161 | 162 | if _, err = tx.Stmt(s.stmts["insertPeerTorrent"]).Exec(peerID, torrentID); err != nil { 163 | return fmt.Errorf("savePeer: %s", err) 164 | } 165 | return tx.Commit() 166 | } 167 | 168 | func (s *Store) RemovePeer(p *models.Peer) (err error) { 169 | s.lock.Lock() 170 | defer s.lock.Unlock() 171 | 172 | _, err = s.stmts["removePeer"].Exec(p.Addr.String()) 173 | return err 174 | } 175 | 176 | // TorrentsByHash implements torrentStore 177 | func (s *Store) TorrentByHash(ih models.Infohash) (*models.Torrent, error) { 178 | s.lock.RLock() 179 | defer s.lock.RUnlock() 180 | 181 | rows, err := s.stmts["getTorrent"].Query(ih) 182 | if err != nil { 183 | return nil, err 184 | } 185 | defer rows.Close() 186 | torrents, err := s.fetchTorrents(rows) 187 | if err != nil { 188 | return nil, err 189 | } 190 | return torrents[0], nil 191 | } 192 | 193 | // TorrentsByName implements torrentStore 194 | func (s *Store) TorrentsByName(query string, offset int) ([]*models.Torrent, error) { 195 | s.lock.RLock() 196 | defer s.lock.RUnlock() 197 | 198 | rows, err := s.stmts["searchTorrents"].Query(fmt.Sprintf("%%%s%%", query), offset) 199 | if err != nil { 200 | return nil, err 201 | } 202 | defer rows.Close() 203 | torrents, err := s.fetchTorrents(rows) 204 | if err != nil { 205 | return nil, err 206 | } 207 | return torrents, nil 208 | } 209 | 210 | // TorrentsByTag implements torrentStore 211 | func (s *Store) TorrentsByTag(tag string, offset int) ([]*models.Torrent, error) { 212 | s.lock.RLock() 213 | defer s.lock.RUnlock() 214 | 215 | rows, err := s.stmts["torrentsByTag"].Query(tag, offset) 216 | if err != nil { 217 | return nil, err 218 | } 219 | defer rows.Close() 220 | torrents, err := s.fetchTorrents(rows) 221 | if err != nil { 222 | return nil, err 223 | } 224 | return torrents, nil 225 | } 226 | 227 | // SaveTag implements tagStore interface 228 | func (s *Store) SaveTag(tag string) (int, error) { 229 | s.lock.Lock() 230 | defer s.lock.Unlock() 231 | 232 | res, err := s.stmts["insertTag"].Exec(tag) 233 | if err != nil { 234 | return 0, fmt.Errorf("saveTag: %s", err) 235 | } 236 | tagID, err := res.LastInsertId() 237 | if err != nil { 238 | return 0, fmt.Errorf("saveTag: %s", err) 239 | } 240 | return int(tagID), nil 241 | } 242 | 243 | func (s *Store) fetchTorrents(rows *sql.Rows) (torrents []*models.Torrent, err error) { 244 | for rows.Next() { 245 | var t models.Torrent 246 | /* 247 | t := &models.Torrent{ 248 | Files: []models.File{}, 249 | Tags: []string{}, 250 | } 251 | */ 252 | err = rows.Scan( 253 | &t.ID, &t.Infohash, &t.Name, &t.Size, &t.Created, &t.Updated, 254 | ) 255 | if err != nil { 256 | return nil, err 257 | } 258 | 259 | err = func() error { 260 | rowsf, err := s.stmts["selectFiles"].Query(t.ID) 261 | defer rowsf.Close() 262 | if err != nil { 263 | return fmt.Errorf("failed to select files: %s", err) 264 | } 265 | for rowsf.Next() { 266 | var f models.File 267 | err = rowsf.Scan(&f.ID, &f.TorrentID, &f.Path, &f.Size) 268 | if err != nil { 269 | return fmt.Errorf("failed to build file: %s", err) 270 | } 271 | } 272 | return nil 273 | }() 274 | if err != nil { 275 | return nil, err 276 | } 277 | 278 | err = func() error { 279 | rowst, err := s.stmts["selectTags"].Query(t.ID) 280 | defer rowst.Close() 281 | if err != nil { 282 | return fmt.Errorf("failed to select tags: %s", err) 283 | } 284 | for rowst.Next() { 285 | var tg string 286 | err = rowst.Scan(&tg) 287 | if err != nil { 288 | return fmt.Errorf("failed to build tag: %s", err) 289 | } 290 | t.Tags = append(t.Tags, tg) 291 | } 292 | return nil 293 | }() 294 | if err != nil { 295 | return nil, err 296 | } 297 | torrents = append(torrents, &t) 298 | } 299 | return torrents, err 300 | } 301 | 302 | func (s *Store) migrate() error { 303 | _, err := s.conn.Exec(` 304 | pragma journal_mode=wal; 305 | pragma temp_store=1; 306 | pragma foreign_keys=on; 307 | pragma encoding='utf-8'; 308 | `) 309 | if err != nil { 310 | return err 311 | } 312 | 313 | tx, err := s.conn.Begin() 314 | if err != nil { 315 | return err 316 | } 317 | defer tx.Rollback() 318 | 319 | var version int 320 | err = tx.QueryRow("pragma user_version;").Scan(&version) 321 | if err != nil { 322 | return err 323 | } 324 | 325 | if version == 0 { 326 | _, err = tx.Exec(sqliteSchema) 327 | if err != nil { 328 | return err 329 | } 330 | } 331 | tx.Commit() 332 | 333 | return nil 334 | } 335 | 336 | func (s *Store) prepareStatements() error { 337 | var err error 338 | if s.stmts["removeTorrent"], err = s.conn.Prepare( 339 | `delete from torrents 340 | where infohash = ?`, 341 | ); err != nil { 342 | return err 343 | } 344 | 345 | if s.stmts["selectPendingInfohashes"], err = s.conn.Prepare( 346 | `select max(p.address) as address, t.infohash 347 | from torrents t 348 | join peers_torrents pt on pt.torrent_id = t.id 349 | join peers p on p.id = pt.peer_id 350 | where t.name is null 351 | group by t.infohash 352 | limit ?`, 353 | ); err != nil { 354 | return err 355 | } 356 | 357 | if s.stmts["selectFiles"], err = s.conn.Prepare( 358 | `select * from files 359 | where torrent_id = ? 360 | order by path asc`, 361 | ); err != nil { 362 | return err 363 | } 364 | 365 | if s.stmts["insertPeer"], err = s.conn.Prepare( 366 | `insert or ignore into peers 367 | (address, created, updated) 368 | values 369 | (?, date('now'), date('now'))`, 370 | ); err != nil { 371 | return err 372 | } 373 | 374 | if s.stmts["insertPeerTorrent"], err = s.conn.Prepare( 375 | `insert or ignore into peers_torrents 376 | (peer_id, torrent_id) 377 | values 378 | (?, ?)`, 379 | ); err != nil { 380 | return err 381 | } 382 | 383 | if s.stmts["insertTorrent"], err = s.conn.Prepare( 384 | `insert or replace into torrents ( 385 | name, infohash, size, created, updated 386 | ) values ( 387 | ?, ?, ?, date('now'), date('now') 388 | )`, 389 | ); err != nil { 390 | return err 391 | } 392 | 393 | if s.stmts["getTorrent"], err = s.conn.Prepare( 394 | `select * from torrents where infohash = ? limit 1`, 395 | ); err != nil { 396 | return err 397 | } 398 | 399 | if s.stmts["insertFile"], err = s.conn.Prepare( 400 | `insert into files 401 | (torrent_id, path, size) 402 | values 403 | (?, ?, ?)`, 404 | ); err != nil { 405 | return err 406 | } 407 | 408 | if s.stmts["selectTags"], err = s.conn.Prepare( 409 | `select name 410 | from tags t 411 | inner join tags_torrents tt on t.id = tt.tag_id 412 | where tt.torrent_id = ?`, 413 | ); err != nil { 414 | return err 415 | } 416 | 417 | if s.stmts["removePeer"], err = s.conn.Prepare( 418 | `delete from peers where address = ?`, 419 | ); err != nil { 420 | return err 421 | } 422 | 423 | if s.stmts["insertTagTorrent"], err = s.conn.Prepare( 424 | `insert or ignore into tags_torrents 425 | (tag_id, torrent_id) values (?, ?)`, 426 | ); err != nil { 427 | return err 428 | } 429 | 430 | if s.stmts["insertTag"], err = s.conn.Prepare( 431 | `insert or replace into tags (name) values (?)`, 432 | ); err != nil { 433 | return err 434 | } 435 | 436 | if s.stmts["torrentsByTag"], err = s.conn.Prepare( 437 | `select t.id, t.infohash, t.name, t.size, t.created, t.updated 438 | from torrents t 439 | inner join tags_torrents tt on t.id = tt.torrent_id 440 | inner join tags ta on tt.tag_id = ta.id 441 | where ta.name = ? group by t.id 442 | order by updated asc 443 | limit 50 offset ?`, 444 | ); err != nil { 445 | return err 446 | } 447 | 448 | if s.stmts["searchTorrents"], err = s.conn.Prepare( 449 | `select id, infohash, name, size, updated 450 | from torrents 451 | where id in ( 452 | select * from torrents_fts 453 | where torrents_fts match ? 454 | order by rank desc 455 | ) 456 | order by updated desc 457 | limit 50 offset ?`, 458 | ); err != nil { 459 | return err 460 | } 461 | 462 | return nil 463 | } 464 | 465 | const sqliteSchema = `create table if not exists torrents ( 466 | id integer primary key, 467 | infohash blob not null unique, 468 | size bigint, 469 | name text, 470 | created timestamp with time zone, 471 | updated timestamp with time zone, 472 | tsv tsvector 473 | ); 474 | create virtual table torrents_fts using fts5( 475 | name, content='torrents', content_rowid='id', 476 | tokenize="porter unicode61 separators ' !""#$%&''()*+,-./:;<=>?@[\]^_` + "`" + `{|}~'" 477 | ); 478 | create trigger torrents_after_insert after insert on torrents begin 479 | insert into torrents_fts(rowid, name) values (new.id, new.name); 480 | end; 481 | create trigger torrents_ad after delete on torrents begin 482 | insert into torrents_fts(torrents_fts, rowid, name) values('delete', old.id, old.name); 483 | end; 484 | create trigger torrents_au after update on torrents begin 485 | insert into torrents_fts(torrents_fts, rowid, name) values('delete', old.id, old.name); 486 | insert into torrents_fts(rowid, name) values (new.id, new.name); 487 | end; 488 | create table if not exists files ( 489 | id integer primary key, 490 | torrent_id integer not null references torrents on delete cascade, 491 | path text, 492 | size bigint 493 | ); 494 | create index files_torrent_idx on files (torrent_id); 495 | create table if not exists tags ( 496 | id integer primary key, 497 | name character varying(50) unique 498 | ); 499 | create unique index tags_name_idx on tags (name); 500 | create table if not exists tags_torrents ( 501 | tag_id integer not null references tags on delete cascade, 502 | torrent_id integer not null references torrents on delete cascade, 503 | primary key (tag_id, torrent_id) 504 | ); 505 | create index tags_torrents_tag_idx on tags_torrents (tag_id); 506 | create index tags_torrents_torrent_idx on tags_torrents (torrent_id); 507 | create table if not exists peers ( 508 | id integer primary key, 509 | address character varying(50) not null unique, 510 | created timestamp with time zone, 511 | updated timestamp with time zone 512 | ); 513 | create table if not exists peers_torrents ( 514 | peer_id integer not null references peers on delete cascade, 515 | torrent_id integer not null references torrents on delete cascade, 516 | primary key (peer_id, torrent_id) 517 | ); 518 | create index peers_torrents_peer_idx on peers_torrents (peer_id); 519 | create index peers_torrents_torrent_idx on peers_torrents (torrent_id); 520 | pragma user_version = 1;` 521 | -------------------------------------------------------------------------------- /dht/messages.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | import ( 4 | "fmt" 5 | "net" 6 | 7 | "src.userspace.com.au/dhtsearch/krpc" 8 | "src.userspace.com.au/dhtsearch/models" 9 | ) 10 | 11 | func (n *Node) onPingQuery(rn remoteNode, msg map[string]interface{}) error { 12 | t, err := krpc.GetString(msg, "t") 13 | if err != nil { 14 | return err 15 | } 16 | n.queueMsg(rn, krpc.MakeResponse(t, map[string]interface{}{ 17 | "id": string(n.id), 18 | })) 19 | return nil 20 | } 21 | 22 | func (n *Node) onGetPeersQuery(rn remoteNode, msg map[string]interface{}) error { 23 | a, err := krpc.GetMap(msg, "a") 24 | if err != nil { 25 | return err 26 | } 27 | 28 | // This is the ih of the torrent 29 | torrent, err := krpc.GetString(a, "info_hash") 30 | if err != nil { 31 | return err 32 | } 33 | th, err := models.InfohashFromString(torrent) 34 | if err != nil { 35 | return err 36 | } 37 | //n.log.Debug("get_peers query", "source", rn, "torrent", th) 38 | 39 | token := torrent[:2] 40 | neighbour := models.GenerateNeighbour(n.id, *th) 41 | /* 42 | nodes := n.rTable.get(8) 43 | compactNS := []string{} 44 | for _, rn := range nodes { 45 | ns := encodeCompactNodeAddr(rn.addr.String()) 46 | if ns == "" { 47 | n.log.Warn("failed to compact node", "address", rn.address.String()) 48 | continue 49 | } 50 | compactNS = append(compactNS, ns) 51 | } 52 | */ 53 | 54 | t := msg["t"].(string) 55 | n.queueMsg(rn, krpc.MakeResponse(t, map[string]interface{}{ 56 | "id": string(neighbour), 57 | "token": token, 58 | "nodes": "", 59 | //"nodes": strings.Join(compactNS, ""), 60 | })) 61 | 62 | //nodes := n.rTable.get(50) 63 | /* 64 | fmt.Printf("sending get_peers for %s to %d nodes\n", *th, len(nodes)) 65 | q := krpc.MakeQuery(newTransactionID(), "get_peers", map[string]interface{}{ 66 | "id": string(id), 67 | "info_hash": string(*th), 68 | }) 69 | for _, o := range nodes { 70 | n.queueMsg(*o, q) 71 | } 72 | */ 73 | return nil 74 | } 75 | 76 | func (n *Node) onAnnouncePeerQuery(rn remoteNode, msg map[string]interface{}) error { 77 | a, err := krpc.GetMap(msg, "a") 78 | if err != nil { 79 | return err 80 | } 81 | 82 | n.log.Debug("announce_peer", "source", rn) 83 | 84 | host, port, err := net.SplitHostPort(rn.addr.String()) 85 | if err != nil { 86 | return err 87 | } 88 | if port == "0" { 89 | return fmt.Errorf("ignoring port 0") 90 | } 91 | 92 | ihStr, err := krpc.GetString(a, "info_hash") 93 | if err != nil { 94 | return err 95 | } 96 | ih, err := models.InfohashFromString(ihStr) 97 | if err != nil { 98 | return fmt.Errorf("invalid torrent: %s", err) 99 | } 100 | 101 | newPort, err := krpc.GetInt(a, "port") 102 | if err == nil { 103 | if iPort, err := krpc.GetInt(a, "implied_port"); err == nil && iPort == 0 { 104 | // Use the port in the message 105 | addr, err := net.ResolveUDPAddr(n.family, fmt.Sprintf("%s:%d", host, newPort)) 106 | if err != nil { 107 | return err 108 | } 109 | n.log.Debug("implied port", "infohash", ih, "original", rn.addr.String(), "new", addr.String()) 110 | rn = remoteNode{addr: addr, id: rn.id} 111 | } 112 | } 113 | 114 | // TODO do we reply? 115 | 116 | p := models.Peer{Addr: rn.addr, Infohash: *ih} 117 | if n.OnAnnouncePeer != nil { 118 | go n.OnAnnouncePeer(p) 119 | } 120 | return nil 121 | } 122 | 123 | func (n *Node) onFindNodeResponse(rn remoteNode, msg map[string]interface{}) { 124 | r := msg["r"].(map[string]interface{}) 125 | nodes := r["nodes"].(string) 126 | n.processFindNodeResults(rn, nodes) 127 | } 128 | -------------------------------------------------------------------------------- /dht/node.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "net" 7 | "time" 8 | 9 | "github.com/hashicorp/golang-lru" 10 | "golang.org/x/time/rate" 11 | "src.userspace.com.au/dhtsearch/krpc" 12 | "src.userspace.com.au/dhtsearch/models" 13 | "src.userspace.com.au/go-bencode" 14 | "src.userspace.com.au/logger" 15 | ) 16 | 17 | var ( 18 | routers = []string{ 19 | "dht.libtorrent.org:25401", 20 | "router.bittorrent.com:6881", 21 | "dht.transmissionbt.com:6881", 22 | "router.utorrent.com:6881", 23 | "dht.aelitis.com:6881", 24 | } 25 | ) 26 | 27 | // Node joins the DHT network 28 | type Node struct { 29 | id models.Infohash 30 | family string 31 | address string 32 | port int 33 | conn net.PacketConn 34 | pool chan chan packet 35 | rTable *routingTable 36 | udpTimeout int 37 | packetsOut chan packet 38 | log logger.Logger 39 | limiter *rate.Limiter 40 | blacklist *lru.ARCCache 41 | 42 | // OnAnnoucePeer is called for each peer that announces itself 43 | OnAnnouncePeer func(models.Peer) 44 | // OnBadPeer is called for each bad peer 45 | OnBadPeer func(models.Peer) 46 | } 47 | 48 | // NewNode creates a new DHT node 49 | func NewNode(opts ...Option) (*Node, error) { 50 | var err error 51 | id := models.GenInfohash() 52 | 53 | n := &Node{ 54 | id: id, 55 | family: "udp4", 56 | port: 6881, 57 | udpTimeout: 10, 58 | limiter: rate.NewLimiter(rate.Limit(100000), 2000000), 59 | log: logger.New(&logger.Options{Name: "dht"}), 60 | } 61 | 62 | n.rTable, err = newRoutingTable(id, 2000) 63 | if err != nil { 64 | n.log.Error("failed to create routing table", "error", err) 65 | return nil, err 66 | } 67 | 68 | // Set variadic options passed 69 | for _, option := range opts { 70 | err = option(n) 71 | if err != nil { 72 | return nil, err 73 | } 74 | } 75 | 76 | if n.blacklist == nil { 77 | n.blacklist, err = lru.NewARC(1000) 78 | if err != nil { 79 | return nil, err 80 | } 81 | } 82 | 83 | if n.family != "udp4" { 84 | n.log.Debug("trying udp6 server") 85 | n.conn, err = net.ListenPacket("udp6", fmt.Sprintf("[%s]:%d", net.IPv6zero.String(), n.port)) 86 | if err == nil { 87 | n.family = "udp6" 88 | } 89 | } 90 | if n.conn == nil { 91 | n.conn, err = net.ListenPacket("udp4", fmt.Sprintf("%s:%d", net.IPv4zero.String(), n.port)) 92 | if err == nil { 93 | n.family = "udp4" 94 | } 95 | } 96 | if err != nil { 97 | n.log.Error("failed to listen", "error", err) 98 | return nil, err 99 | } 100 | n.log.Info("listening", "id", n.id, "network", n.family, "address", n.conn.LocalAddr().String()) 101 | 102 | return n, nil 103 | } 104 | 105 | // Close stuff 106 | func (n *Node) Close() error { 107 | n.log.Warn("node closing") 108 | return nil 109 | } 110 | 111 | // Run starts the node on the DHT 112 | func (n *Node) Run() { 113 | // Packets onto the network 114 | n.packetsOut = make(chan packet, 1024) 115 | 116 | // Create a slab for allocation 117 | byteSlab := newSlab(8192, 10) 118 | 119 | n.log.Debug("starting packet writer") 120 | go n.packetWriter() 121 | 122 | // Find neighbours 123 | go n.makeNeighbours() 124 | 125 | n.log.Debug("starting packet reader") 126 | for { 127 | b := byteSlab.alloc() 128 | c, addr, err := n.conn.ReadFrom(b) 129 | if err != nil { 130 | n.log.Warn("UDP read error", "error", err) 131 | return 132 | } 133 | 134 | // Chop and process 135 | n.processPacket(packet{ 136 | data: b[0:c], 137 | raddr: addr, 138 | }) 139 | byteSlab.free(b) 140 | } 141 | } 142 | 143 | func (n *Node) makeNeighbours() { 144 | // TODO configurable 145 | ticker := time.Tick(5 * time.Second) 146 | 147 | n.bootstrap() 148 | 149 | for { 150 | select { 151 | case <-ticker: 152 | if n.rTable.isEmpty() { 153 | n.bootstrap() 154 | } else { 155 | // Send to all nodes 156 | nodes := n.rTable.get(0) 157 | for _, rn := range nodes { 158 | n.findNode(rn, models.GenerateNeighbour(n.id, rn.id)) 159 | } 160 | n.rTable.flush() 161 | } 162 | } 163 | } 164 | } 165 | 166 | func (n *Node) bootstrap() { 167 | n.log.Debug("bootstrapping") 168 | for _, s := range routers { 169 | addr, err := net.ResolveUDPAddr(n.family, s) 170 | if err != nil { 171 | n.log.Error("failed to parse bootstrap address", "error", err) 172 | continue 173 | } 174 | rn := &remoteNode{addr: addr} 175 | n.findNode(rn, n.id) 176 | } 177 | } 178 | 179 | func (n *Node) packetWriter() { 180 | for p := range n.packetsOut { 181 | if p.raddr.String() == n.conn.LocalAddr().String() { 182 | continue 183 | } 184 | ctx, cancel := context.WithCancel(context.Background()) 185 | defer cancel() 186 | if err := n.limiter.WaitN(ctx, len(p.data)); err != nil { 187 | n.log.Warn("rate limited", "error", err) 188 | continue 189 | } 190 | //n.log.Debug("writing packet", "dest", p.raddr.String()) 191 | _, err := n.conn.WriteTo(p.data, p.raddr) 192 | if err != nil { 193 | n.blacklist.Add(p.raddr.String(), true) 194 | // TODO reduce limit 195 | n.log.Warn("failed to write packet", "error", err) 196 | if n.OnBadPeer != nil { 197 | peer := models.Peer{Addr: p.raddr} 198 | go n.OnBadPeer(peer) 199 | } 200 | } 201 | } 202 | } 203 | 204 | func (n *Node) findNode(rn *remoteNode, id models.Infohash) { 205 | target := models.GenInfohash() 206 | n.sendQuery(rn, "find_node", map[string]interface{}{ 207 | "id": string(id), 208 | "target": string(target), 209 | }) 210 | } 211 | 212 | // ping sends ping query to the chan. 213 | func (n *Node) ping(rn *remoteNode) { 214 | id := models.GenerateNeighbour(n.id, rn.id) 215 | n.sendQuery(rn, "ping", map[string]interface{}{ 216 | "id": string(id), 217 | }) 218 | } 219 | 220 | func (n *Node) sendQuery(rn *remoteNode, qType string, a map[string]interface{}) error { 221 | // Stop if sending to self 222 | if rn.id.Equal(n.id) { 223 | return nil 224 | } 225 | 226 | t := krpc.NewTransactionID() 227 | 228 | data := krpc.MakeQuery(t, qType, a) 229 | b, err := bencode.Encode(data) 230 | if err != nil { 231 | return err 232 | } 233 | //fmt.Printf("sending %s to %s\n", qType, rn.String()) 234 | n.packetsOut <- packet{ 235 | data: b, 236 | raddr: rn.addr, 237 | } 238 | return nil 239 | } 240 | 241 | // Parse a KRPC packet into a message 242 | func (n *Node) processPacket(p packet) error { 243 | response, _, err := bencode.DecodeDict(p.data, 0) 244 | if err != nil { 245 | return err 246 | } 247 | 248 | y, err := krpc.GetString(response, "y") 249 | if err != nil { 250 | return err 251 | } 252 | 253 | if _, black := n.blacklist.Get(p.raddr.String()); black { 254 | return fmt.Errorf("blacklisted: %s", p.raddr.String()) 255 | } 256 | 257 | switch y { 258 | case "q": 259 | err = n.handleRequest(p.raddr, response) 260 | case "r": 261 | err = n.handleResponse(p.raddr, response) 262 | case "e": 263 | err = n.handleError(p.raddr, response) 264 | default: 265 | err = fmt.Errorf("missing request type") 266 | } 267 | if err != nil { 268 | n.log.Warn("failed to process packet", "error", err) 269 | n.blacklist.Add(p.raddr.String(), true) 270 | } 271 | return err 272 | } 273 | 274 | // bencode data and send 275 | func (n *Node) queueMsg(rn remoteNode, data map[string]interface{}) error { 276 | b, err := bencode.Encode(data) 277 | if err != nil { 278 | return err 279 | } 280 | n.packetsOut <- packet{ 281 | data: b, 282 | raddr: rn.addr, 283 | } 284 | return nil 285 | } 286 | 287 | // handleRequest handles the requests received from udp. 288 | func (n *Node) handleRequest(addr net.Addr, m map[string]interface{}) error { 289 | q, err := krpc.GetString(m, "q") 290 | if err != nil { 291 | return err 292 | } 293 | 294 | a, err := krpc.GetMap(m, "a") 295 | if err != nil { 296 | return err 297 | } 298 | 299 | id, err := krpc.GetString(a, "id") 300 | if err != nil { 301 | return err 302 | } 303 | 304 | ih, err := models.InfohashFromString(id) 305 | if err != nil { 306 | return err 307 | } 308 | 309 | if n.id.Equal(*ih) { 310 | return nil 311 | } 312 | 313 | rn := &remoteNode{addr: addr, id: *ih} 314 | 315 | switch q { 316 | case "ping": 317 | err = n.onPingQuery(*rn, m) 318 | 319 | case "get_peers": 320 | err = n.onGetPeersQuery(*rn, m) 321 | 322 | case "announce_peer": 323 | n.onAnnouncePeerQuery(*rn, m) 324 | 325 | default: 326 | //n.queueMsg(addr, makeError(t, protocolError, "invalid q")) 327 | return nil 328 | } 329 | n.rTable.add(rn) 330 | return err 331 | } 332 | 333 | // handleResponse handles responses received from udp. 334 | func (n *Node) handleResponse(addr net.Addr, m map[string]interface{}) error { 335 | r, err := krpc.GetMap(m, "r") 336 | if err != nil { 337 | return err 338 | } 339 | id, err := krpc.GetString(r, "id") 340 | if err != nil { 341 | return err 342 | } 343 | ih, err := models.InfohashFromString(id) 344 | if err != nil { 345 | return err 346 | } 347 | 348 | rn := &remoteNode{addr: addr, id: *ih} 349 | 350 | nodes, err := krpc.GetString(r, "nodes") 351 | // find_nodes/get_peers response with nodes 352 | if err == nil { 353 | n.onFindNodeResponse(*rn, m) 354 | n.processFindNodeResults(*rn, nodes) 355 | n.rTable.add(rn) 356 | return nil 357 | } 358 | 359 | values, err := krpc.GetList(r, "values") 360 | // get_peers response 361 | if err == nil { 362 | n.log.Debug("get_peers response", "source", rn) 363 | for _, v := range values { 364 | addr := krpc.DecodeCompactNodeAddr(v.(string)) 365 | n.log.Debug("unhandled get_peer request", "addres", addr) 366 | 367 | // TODO new peer needs to be matched to previous get_peers request 368 | // n.peersManager.Insert(ih, p) 369 | } 370 | n.rTable.add(rn) 371 | } 372 | return nil 373 | } 374 | 375 | // handleError handles errors received from udp. 376 | func (n *Node) handleError(addr net.Addr, m map[string]interface{}) error { 377 | e, err := krpc.GetList(m, "e") 378 | if err != nil { 379 | return err 380 | } 381 | 382 | if len(e) != 2 { 383 | return fmt.Errorf("error packet wrong length %d", len(e)) 384 | } 385 | code := e[0].(int64) 386 | msg := e[1].(string) 387 | n.log.Debug("error packet", "address", addr.String(), "code", code, "error", msg) 388 | 389 | return nil 390 | } 391 | 392 | // Process another node's response to a find_node query. 393 | func (n *Node) processFindNodeResults(rn remoteNode, nodeList string) { 394 | nodeLength := krpc.IPv4NodeAddrLen 395 | if n.family == "udp6" { 396 | nodeLength = krpc.IPv6NodeAddrLen 397 | } 398 | 399 | if len(nodeList)%nodeLength != 0 { 400 | n.log.Error("node list is wrong length", "length", len(nodeList)) 401 | n.blacklist.Add(rn.addr.String(), true) 402 | return 403 | } 404 | 405 | //fmt.Printf("%s sent %d nodes\n", rn.address.String(), len(nodeList)/nodeLength) 406 | 407 | // We got a byte array in groups of 26 or 38 408 | for i := 0; i < len(nodeList); i += nodeLength { 409 | id := nodeList[i : i+models.InfohashLength] 410 | addrStr := krpc.DecodeCompactNodeAddr(nodeList[i+models.InfohashLength : i+nodeLength]) 411 | 412 | ih, err := models.InfohashFromString(id) 413 | if err != nil { 414 | n.log.Warn("invalid infohash in node list") 415 | continue 416 | } 417 | 418 | addr, err := net.ResolveUDPAddr(n.family, addrStr) 419 | if err != nil || addr.Port == 0 { 420 | //n.log.Warn("unable to resolve", "address", addrStr, "error", err) 421 | continue 422 | } 423 | 424 | rn := &remoteNode{addr: addr, id: *ih} 425 | n.rTable.add(rn) 426 | } 427 | } 428 | -------------------------------------------------------------------------------- /dht/options.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | import ( 4 | "github.com/hashicorp/golang-lru" 5 | "src.userspace.com.au/dhtsearch/models" 6 | "src.userspace.com.au/logger" 7 | ) 8 | 9 | type Option func(*Node) error 10 | 11 | func SetOnAnnouncePeer(f func(models.Peer)) Option { 12 | return func(n *Node) error { 13 | n.OnAnnouncePeer = f 14 | return nil 15 | } 16 | } 17 | 18 | func SetOnBadPeer(f func(models.Peer)) Option { 19 | return func(n *Node) error { 20 | n.OnBadPeer = f 21 | return nil 22 | } 23 | } 24 | 25 | // SetAddress sets the IP address to listen on 26 | func SetAddress(ip string) Option { 27 | return func(n *Node) error { 28 | n.address = ip 29 | return nil 30 | } 31 | } 32 | 33 | // SetPort sets the port to listen on 34 | func SetPort(p int) Option { 35 | return func(n *Node) error { 36 | n.port = p 37 | return nil 38 | } 39 | } 40 | 41 | // SetIPv6 enables IPv6 42 | func SetIPv6(b bool) Option { 43 | return func(n *Node) error { 44 | if b { 45 | n.family = "udp6" 46 | } 47 | return nil 48 | } 49 | } 50 | 51 | // SetUDPTimeout sets the number of seconds to wait for UDP connections 52 | func SetUDPTimeout(s int) Option { 53 | return func(n *Node) error { 54 | n.udpTimeout = s 55 | return nil 56 | } 57 | } 58 | 59 | // SetLogger sets the logger 60 | func SetLogger(l logger.Logger) Option { 61 | return func(n *Node) error { 62 | n.log = l 63 | return nil 64 | } 65 | } 66 | 67 | // SetBlacklist sets the size of the node blacklist 68 | func SetBlacklist(bl *lru.ARCCache) Option { 69 | return func(n *Node) (err error) { 70 | n.blacklist = bl 71 | return err 72 | } 73 | } 74 | -------------------------------------------------------------------------------- /dht/packet.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | import "net" 4 | 5 | // Arbitrary packet types 6 | // Order these lowest to highest priority for use in 7 | // priority queue heap 8 | const ( 9 | _ int = iota 10 | pktQPing 11 | pktRPing 12 | pktQFindNode 13 | pktRAnnouncePeer 14 | pktRGetPeers 15 | ) 16 | 17 | var pktName = map[int]string{ 18 | pktQFindNode: "find_node", 19 | pktQPing: "ping", 20 | pktRPing: "ping", 21 | pktRAnnouncePeer: "annouce_peer", 22 | pktRGetPeers: "get_peers", 23 | } 24 | 25 | // Unprocessed packet from socket 26 | type packet struct { 27 | // The packet type 28 | //priority int 29 | // Required by heap interface 30 | //index int 31 | data []byte 32 | raddr net.Addr 33 | } 34 | -------------------------------------------------------------------------------- /dht/remote_node.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | import ( 4 | "fmt" 5 | "net" 6 | 7 | "src.userspace.com.au/dhtsearch/models" 8 | ) 9 | 10 | type remoteNode struct { 11 | addr net.Addr 12 | id models.Infohash 13 | } 14 | 15 | // String implements fmt.Stringer 16 | func (r remoteNode) String() string { 17 | return fmt.Sprintf("%s (%s)", r.id.String(), r.addr.String()) 18 | } 19 | -------------------------------------------------------------------------------- /dht/routing_table.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | import ( 4 | "container/heap" 5 | "sync" 6 | 7 | "src.userspace.com.au/dhtsearch/models" 8 | ) 9 | 10 | type rItem struct { 11 | value *remoteNode 12 | distance int 13 | index int // Index in heap 14 | } 15 | 16 | type priorityQueue []*rItem 17 | 18 | type routingTable struct { 19 | id models.Infohash 20 | max int 21 | items priorityQueue 22 | addresses map[string]*remoteNode 23 | sync.Mutex 24 | } 25 | 26 | func newRoutingTable(id models.Infohash, max int) (*routingTable, error) { 27 | k := &routingTable{ 28 | id: id, 29 | max: max, 30 | } 31 | k.flush() 32 | heap.Init(&k.items) 33 | return k, nil 34 | } 35 | 36 | // Len implements sort.Interface 37 | func (pq priorityQueue) Len() int { return len(pq) } 38 | 39 | // Less implements sort.Interface 40 | func (pq priorityQueue) Less(i, j int) bool { 41 | return pq[i].distance > pq[j].distance 42 | } 43 | 44 | // Swap implements sort.Interface 45 | func (pq priorityQueue) Swap(i, j int) { 46 | pq[i], pq[j] = pq[j], pq[i] 47 | pq[i].index = i 48 | pq[j].index = j 49 | } 50 | 51 | // Push implements heap.Interface 52 | func (pq *priorityQueue) Push(x interface{}) { 53 | n := len(*pq) 54 | item := x.(*rItem) 55 | item.index = n 56 | *pq = append(*pq, item) 57 | } 58 | 59 | // Pop implements heap.Interface 60 | func (pq *priorityQueue) Pop() interface{} { 61 | old := *pq 62 | n := len(old) 63 | item := old[n-1] 64 | item.index = -1 // for safety 65 | *pq = old[0 : n-1] 66 | return item 67 | } 68 | 69 | func (k *routingTable) add(rn *remoteNode) { 70 | // Check IP and ports are valid and not self 71 | if !rn.id.Valid() || rn.id.Equal(k.id) { 72 | return 73 | } 74 | 75 | k.Lock() 76 | defer k.Unlock() 77 | 78 | if _, ok := k.addresses[rn.addr.String()]; ok { 79 | return 80 | } 81 | k.addresses[rn.addr.String()] = rn 82 | 83 | item := &rItem{ 84 | value: rn, 85 | distance: k.id.Distance(rn.id), 86 | } 87 | 88 | heap.Push(&k.items, item) 89 | 90 | if len(k.items) > k.max { 91 | for i := k.max - 1; i < len(k.items); i++ { 92 | old := k.items[i] 93 | delete(k.addresses, old.value.addr.String()) 94 | heap.Remove(&k.items, i) 95 | } 96 | } 97 | } 98 | 99 | func (k *routingTable) get(n int) (out []*remoteNode) { 100 | if n == 0 { 101 | n = len(k.items) 102 | } 103 | for i := 0; i < n && i < len(k.items); i++ { 104 | out = append(out, k.items[i].value) 105 | } 106 | return out 107 | } 108 | 109 | func (k *routingTable) flush() { 110 | k.Lock() 111 | defer k.Unlock() 112 | 113 | k.items = make(priorityQueue, 0) 114 | k.addresses = make(map[string]*remoteNode, k.max) 115 | } 116 | 117 | func (k *routingTable) isEmpty() bool { 118 | k.Lock() 119 | defer k.Unlock() 120 | return len(k.items) == 0 121 | } 122 | -------------------------------------------------------------------------------- /dht/routing_table_test.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | import ( 4 | "fmt" 5 | "net" 6 | "testing" 7 | 8 | "src.userspace.com.au/dhtsearch/models" 9 | ) 10 | 11 | func TestPriorityQueue(t *testing.T) { 12 | id := "d1c5676ae7ac98e8b19f63565905105e3c4c37a2" 13 | 14 | tests := []string{ 15 | "d1c5676ae7ac98e8b19f63565905105e3c4c37b9", 16 | "d1c5676ae7ac98e8b19f63565905105e3c4c37a9", 17 | "d1c5676ae7ac98e8b19f63565905105e3c4c37a4", 18 | "d1c5676ae7ac98e8b19f63565905105e3c4c37a3", // distance of 159 19 | } 20 | 21 | ih, err := models.InfohashFromString(id) 22 | if err != nil { 23 | t.Errorf("failed to create infohash: %s\n", err) 24 | } 25 | 26 | pq, err := newRoutingTable(*ih, 3) 27 | if err != nil { 28 | t.Errorf("failed to create kTable: %s\n", err) 29 | } 30 | 31 | for i, idt := range tests { 32 | iht, err := models.InfohashFromString(idt) 33 | if err != nil { 34 | t.Errorf("failed to create infohash: %s\n", err) 35 | } 36 | addr, _ := net.ResolveUDPAddr("udp", fmt.Sprintf("0.0.0.0:%d", i)) 37 | pq.add(&remoteNode{id: *iht, addr: addr}) 38 | } 39 | 40 | if len(pq.items) != len(pq.addresses) { 41 | t.Errorf("items and addresses out of sync") 42 | } 43 | 44 | first := pq.items[0].value.id 45 | if first.String() != "d1c5676ae7ac98e8b19f63565905105e3c4c37a3" { 46 | t.Errorf("first is %s with distance %d\n", first, ih.Distance(first)) 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /dht/slab.go: -------------------------------------------------------------------------------- 1 | package dht 2 | 3 | // Slab memory allocation 4 | 5 | // Initialise the slab as a channel of blocks, allocating them as required and 6 | // pushing them back on the slab. This reduces garbage collection. 7 | type slab chan []byte 8 | 9 | func newSlab(blockSize int, numBlocks int) slab { 10 | s := make(slab, numBlocks) 11 | for i := 0; i < numBlocks; i++ { 12 | s <- make([]byte, blockSize) 13 | } 14 | return s 15 | } 16 | 17 | func (s slab) alloc() (x []byte) { 18 | return <-s 19 | } 20 | 21 | func (s slab) free(x []byte) { 22 | // Check we are using the right dimensions 23 | x = x[:cap(x)] 24 | s <- x 25 | } 26 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module src.userspace.com.au/dhtsearch 2 | 3 | go 1.13 4 | 5 | require ( 6 | github.com/cockroachdb/apd v1.1.0 // indirect 7 | github.com/hashicorp/go-version v1.2.1 // indirect 8 | github.com/hashicorp/golang-lru v0.0.0-20180201235237-0fb14efe8c47 9 | github.com/jackc/fake v0.0.0-20150926172116-812a484cc733 // indirect 10 | github.com/jackc/pgx v3.1.0+incompatible 11 | github.com/lib/pq v1.8.0 // indirect 12 | github.com/mattn/go-sqlite3 v1.14.4 13 | github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect 14 | github.com/pkg/errors v0.8.0 // indirect 15 | github.com/satori/go.uuid v1.2.0 // indirect 16 | github.com/shopspring/decimal v1.2.0 // indirect 17 | golang.org/x/net v0.0.0-20180330215511-b68f30494add // indirect 18 | golang.org/x/time v0.0.0-20180314180208-26559e0f760e 19 | gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect 20 | src.userspace.com.au/go-bencode v0.3.1 21 | src.userspace.com.au/logger v0.1.1 22 | ) 23 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ= 2 | github.com/hashicorp/go-version v1.2.1 h1:zEfKbn2+PDgroKdiOzqiE8rsmLqU2uwi5PB5pBJ3TkI= 3 | github.com/hashicorp/go-version v1.2.1/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA= 4 | github.com/hashicorp/golang-lru v0.0.0-20180201235237-0fb14efe8c47 h1:UnszMmmmm5vLwWzDjTFVIkfhvWF1NdrmChl8L2NUDCw= 5 | github.com/hashicorp/golang-lru v0.0.0-20180201235237-0fb14efe8c47/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8= 6 | github.com/jackc/fake v0.0.0-20150926172116-812a484cc733/go.mod h1:WrMFNQdiFJ80sQsxDoMokWK1W5TQtxBFNpzWTD84ibQ= 7 | github.com/jackc/pgx v3.1.0+incompatible/go.mod h1:0ZGrqGqkRlliWnWB4zKnWtjbSWbGkVEFm4TeybAXq+I= 8 | github.com/jackpal/bencode-go v0.0.0-20180813173944-227668e840fa h1:ym9I4Q1lJG8nu+j5R2H6mHOfVjYbSiwUOzh/AFs3Xfs= 9 | github.com/jackpal/bencode-go v0.0.0-20180813173944-227668e840fa/go.mod h1:5FSBQ74yhCl5oQ+QxRPYzWMONFnxbL68/23eezsBI5c= 10 | github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= 11 | github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= 12 | github.com/lib/pq v1.8.0 h1:9xohqzkUwzR4Ga4ivdTcawVS89YSDVxXMa3xJX3cGzg= 13 | github.com/lib/pq v1.8.0/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= 14 | github.com/marksamman/bencode v0.0.0-20150821143521-dc84f26e086e h1:KMs6SK8iDSR1+ZzOK10L5wGPpWDByyvOe5nrqk51g2U= 15 | github.com/marksamman/bencode v0.0.0-20150821143521-dc84f26e086e/go.mod h1:+AHfJo5+69p+fjvMJTmYajNP9rFBHQcaTDFmuXRRATI= 16 | github.com/mattn/go-sqlite3 v1.14.4 h1:4rQjbDxdu9fSgI/r3KN72G3c2goxknAqHHgPWWs8UlI= 17 | github.com/mattn/go-sqlite3 v1.14.4/go.mod h1:WVKg1VTActs4Qso6iwGbiFih2UIHo0ENGwNd0Lj+XmI= 18 | github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs= 19 | github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno= 20 | github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 21 | github.com/satori/go.uuid v1.2.0/go.mod h1:dA0hQrYB0VpLJoorglMZABFdXlWrHn1NEOzdhQKdks0= 22 | github.com/shopspring/decimal v1.2.0 h1:abSATXmQEYyShuxI4/vyW3tV1MrKAJzCZ/0zLUXYbsQ= 23 | github.com/shopspring/decimal v1.2.0/go.mod h1:DKyhrW/HYNuLGql+MJL6WCR6knT2jwCFRcu2hWCYk4o= 24 | golang.org/x/net v0.0.0-20180330215511-b68f30494add h1:oGr9qHpQTQvl/BmeWw95ZrQKahW4qdIPUiGfQkJYDsA= 25 | golang.org/x/net v0.0.0-20180330215511-b68f30494add/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 26 | golang.org/x/time v0.0.0-20180314180208-26559e0f760e h1:aUMCDtB7fbxaw60p2ngy69FCEzU3XpcAEpszqXsdXWg= 27 | golang.org/x/time v0.0.0-20180314180208-26559e0f760e/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= 28 | gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b h1:QRR6H1YWRnHb4Y/HeNFCTJLFVxaq6wH4YuVdsUOr75U= 29 | gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 30 | src.userspace.com.au/go-bencode v0.3.1 h1:O3Qny4bKM4on4U/LjyWb3i7KwQnQnUPIb2PSjGCBnoI= 31 | src.userspace.com.au/go-bencode v0.3.1/go.mod h1:x/2aZW6OnN60Ot7VCTShG1B7H47QaWRS7pCxkgCpfaI= 32 | src.userspace.com.au/logger v0.1.1 h1:mfnQv/pi5X2fxgXxMAJS7oDOWww4p2P2ktDuP5AjCtQ= 33 | src.userspace.com.au/logger v0.1.1/go.mod h1:PC4uwII4fDAhivRMJMp+6uk+VCT6/EZcdDE8hXG56NI= 34 | -------------------------------------------------------------------------------- /krpc/krpc.go: -------------------------------------------------------------------------------- 1 | package krpc 2 | 3 | import ( 4 | "errors" 5 | "fmt" 6 | "math/rand" 7 | "net" 8 | "strconv" 9 | ) 10 | 11 | const ( 12 | transIDBytes = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ" 13 | IPv4NodeAddrLen = 26 14 | IPv6NodeAddrLen = 38 15 | ) 16 | 17 | func NewTransactionID() string { 18 | b := make([]byte, 2) 19 | for i := range b { 20 | b[i] = transIDBytes[rand.Int63()%int64(len(transIDBytes))] 21 | } 22 | return string(b) 23 | } 24 | 25 | // makeQuery returns a query-formed data. 26 | func MakeQuery(transaction, query string, data map[string]interface{}) map[string]interface{} { 27 | return map[string]interface{}{ 28 | "t": transaction, 29 | "y": "q", 30 | "q": query, 31 | "a": data, 32 | } 33 | } 34 | 35 | // makeResponse returns a response-formed data. 36 | func MakeResponse(transaction string, data map[string]interface{}) map[string]interface{} { 37 | return map[string]interface{}{ 38 | "t": transaction, 39 | "y": "r", 40 | "r": data, 41 | } 42 | } 43 | 44 | func GetString(data map[string]interface{}, key string) (string, error) { 45 | val, ok := data[key] 46 | if !ok { 47 | return "", fmt.Errorf("krpc: missing key %s", key) 48 | } 49 | out, ok := val.(string) 50 | if !ok { 51 | return "", fmt.Errorf("krpc: key type mismatch") 52 | } 53 | return out, nil 54 | } 55 | 56 | func GetInt(data map[string]interface{}, key string) (int, error) { 57 | val, ok := data[key] 58 | if !ok { 59 | return 0, fmt.Errorf("krpc: missing key %s", key) 60 | } 61 | out, ok := val.(int64) 62 | if !ok { 63 | return 0, fmt.Errorf("krpc: key type mismatch") 64 | } 65 | return int(out), nil 66 | } 67 | 68 | func GetMap(data map[string]interface{}, key string) (map[string]interface{}, error) { 69 | val, ok := data[key] 70 | if !ok { 71 | return nil, fmt.Errorf("krpc: missing key %s", key) 72 | } 73 | out, ok := val.(map[string]interface{}) 74 | if !ok { 75 | return nil, fmt.Errorf("krpc: key type mismatch") 76 | } 77 | return out, nil 78 | } 79 | 80 | func GetList(data map[string]interface{}, key string) ([]interface{}, error) { 81 | val, ok := data[key] 82 | if !ok { 83 | return nil, fmt.Errorf("krpc: missing key %s", key) 84 | } 85 | out, ok := val.([]interface{}) 86 | if !ok { 87 | return nil, fmt.Errorf("krpc: key type mismatch") 88 | } 89 | return out, nil 90 | } 91 | 92 | // parseKeys parses keys. It just wraps parseKey. 93 | func checkKeys(data map[string]interface{}, pairs [][]string) (err error) { 94 | for _, args := range pairs { 95 | key, t := args[0], args[1] 96 | if err = checkKey(data, key, t); err != nil { 97 | break 98 | } 99 | } 100 | return err 101 | } 102 | 103 | // parseKey parses the key in dict data. `t` is type of the keyed value. 104 | // It's one of "int", "string", "map", "list". 105 | func checkKey(data map[string]interface{}, key string, t string) error { 106 | val, ok := data[key] 107 | if !ok { 108 | return fmt.Errorf("krpc: missing key %s", key) 109 | } 110 | 111 | switch t { 112 | case "string": 113 | _, ok = val.(string) 114 | case "int": 115 | _, ok = val.(int) 116 | case "map": 117 | _, ok = val.(map[string]interface{}) 118 | case "list": 119 | _, ok = val.([]interface{}) 120 | default: 121 | return errors.New("krpc: invalid type") 122 | } 123 | 124 | if !ok { 125 | return errors.New("krpc: key type mismatch") 126 | } 127 | 128 | return nil 129 | } 130 | 131 | // Swiped from nictuku 132 | func DecodeCompactNodeAddr(cni string) string { 133 | if len(cni) == 6 { 134 | return fmt.Sprintf("%d.%d.%d.%d:%d", cni[0], cni[1], cni[2], cni[3], (uint16(cni[4])<<8)|uint16(cni[5])) 135 | } else if len(cni) == 18 { 136 | b := []byte(cni[:16]) 137 | return fmt.Sprintf("[%s]:%d", net.IP.String(b), (uint16(cni[16])<<8)|uint16(cni[17])) 138 | } else { 139 | return "" 140 | } 141 | } 142 | 143 | func EncodeCompactNodeAddr(addr string) string { 144 | var a []uint8 145 | host, port, _ := net.SplitHostPort(addr) 146 | ip := net.ParseIP(host) 147 | if ip == nil { 148 | return "" 149 | } 150 | aa, _ := strconv.ParseUint(port, 10, 16) 151 | c := uint16(aa) 152 | if ip2 := net.IP.To4(ip); ip2 != nil { 153 | a = make([]byte, net.IPv4len+2, net.IPv4len+2) 154 | copy(a, ip2[0:net.IPv4len]) // ignore bytes IPv6 bytes if it's IPv4. 155 | a[4] = byte(c >> 8) 156 | a[5] = byte(c) 157 | } else { 158 | a = make([]byte, net.IPv6len+2, net.IPv6len+2) 159 | copy(a, ip) 160 | a[16] = byte(c >> 8) 161 | a[17] = byte(c) 162 | } 163 | return string(a) 164 | } 165 | 166 | func int2bytes(val int64) []byte { 167 | data, j := make([]byte, 8), -1 168 | for i := 0; i < 8; i++ { 169 | shift := uint64((7 - i) * 8) 170 | data[i] = byte((val & (0xff << shift)) >> shift) 171 | 172 | if j == -1 && data[i] != 0 { 173 | j = i 174 | } 175 | } 176 | 177 | if j != -1 { 178 | return data[j:] 179 | } 180 | return data[:1] 181 | } 182 | -------------------------------------------------------------------------------- /krpc/krpc_test.go: -------------------------------------------------------------------------------- 1 | package krpc 2 | 3 | import ( 4 | "encoding/hex" 5 | "testing" 6 | ) 7 | 8 | func TestCompactNodeAddr(t *testing.T) { 9 | 10 | tests := []struct { 11 | in string 12 | out string 13 | }{ 14 | {in: "192.168.1.1:6881", out: "c0a801011ae1"}, 15 | {in: "[2001:9372:434a:800::2]:6881", out: "20019372434a080000000000000000021ae1"}, 16 | } 17 | 18 | for _, tt := range tests { 19 | r := EncodeCompactNodeAddr(tt.in) 20 | out, _ := hex.DecodeString(tt.out) 21 | if r != string(out) { 22 | t.Errorf("encodeCompactNodeAddr(%s) => %x, expected %s", tt.in, r, tt.out) 23 | } 24 | 25 | s := DecodeCompactNodeAddr(r) 26 | if s != tt.in { 27 | t.Errorf("decodeCompactNodeAddr(%x) => %s, expected %s", r, s, tt.in) 28 | } 29 | } 30 | } 31 | -------------------------------------------------------------------------------- /models/infohash.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import ( 4 | "crypto/sha1" 5 | "encoding/hex" 6 | "fmt" 7 | "io" 8 | "math/rand" 9 | "time" 10 | ) 11 | 12 | const InfohashLength = 20 13 | 14 | // Infohash is a 160 bit (20 byte) value 15 | type Infohash []byte 16 | 17 | // InfohashFromString converts a 40 digit hexadecimal string to an Infohash 18 | func InfohashFromString(s string) (*Infohash, error) { 19 | switch len(s) { 20 | case 20: 21 | // Binary string 22 | ih := Infohash([]byte(s)) 23 | return &ih, nil 24 | case 40: 25 | // Hex string 26 | b, err := hex.DecodeString(s) 27 | if err != nil { 28 | return nil, err 29 | } 30 | ih := Infohash(b) 31 | return &ih, nil 32 | default: 33 | return nil, fmt.Errorf("invalid length %d", len(s)) 34 | } 35 | } 36 | 37 | func (ih Infohash) String() string { 38 | return hex.EncodeToString(ih) 39 | } 40 | 41 | func (ih Infohash) Bytes() []byte { 42 | return []byte(ih) 43 | } 44 | 45 | func (ih Infohash) Valid() bool { 46 | // TODO 47 | return len(ih) == 20 48 | } 49 | 50 | func (ih Infohash) Equal(other Infohash) bool { 51 | if len(ih) != len(other) { 52 | return false 53 | } 54 | for i := 0; i < len(ih); i++ { 55 | if ih[i] != other[i] { 56 | return false 57 | } 58 | } 59 | return true 60 | } 61 | 62 | // Distance determines the distance to another infohash as an integer 63 | func (ih Infohash) Distance(other Infohash) int { 64 | i := 0 65 | for ; i < 20; i++ { 66 | if ih[i] != other[i] { 67 | break 68 | } 69 | } 70 | 71 | if i == 20 { 72 | return 160 73 | } 74 | 75 | xor := ih[i] ^ other[i] 76 | 77 | j := 0 78 | for (xor & 0x80) == 0 { 79 | xor <<= 1 80 | j++ 81 | } 82 | return 8*i + j 83 | } 84 | 85 | func GenerateNeighbour(first, second Infohash) Infohash { 86 | s := append(second[:10], first[10:]...) 87 | return Infohash(s) 88 | } 89 | 90 | func GenInfohash() (ih Infohash) { 91 | random := rand.New(rand.NewSource(time.Now().UnixNano())) 92 | hash := sha1.New() 93 | io.WriteString(hash, time.Now().String()) 94 | io.WriteString(hash, string(random.Int())) 95 | return Infohash(hash.Sum(nil)) 96 | } 97 | -------------------------------------------------------------------------------- /models/infohash_test.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import ( 4 | "encoding/hex" 5 | "testing" 6 | ) 7 | 8 | func TestInfohashImport(t *testing.T) { 9 | 10 | tests := []struct { 11 | str string 12 | ok bool 13 | }{ 14 | {str: "5a3ce1c14e7a08645677bbd1cfe7d8f956d53256", ok: true}, 15 | {str: "5a3ce1c14e7a08645677bbd1cfe7d8f956d53256000", ok: false}, 16 | } 17 | 18 | for _, tt := range tests { 19 | ih, err := InfohashFromString(tt.str) 20 | if tt.ok { 21 | if err != nil { 22 | t.Errorf("FromString failed with %s", err) 23 | } 24 | 25 | idBytes, err := hex.DecodeString(tt.str) 26 | if err != nil { 27 | t.Errorf("failed to decode %s to hex", tt.str) 28 | } 29 | ih2 := Infohash(idBytes) 30 | if !ih.Equal(ih2) { 31 | t.Errorf("expected %s to equal %s", ih, ih2) 32 | } 33 | if ih.String() != tt.str { 34 | t.Errorf("expected ih.String() to equal %s, got %s", tt.str, ih.String()) 35 | } 36 | byt := ih.Bytes() 37 | for i := range byt { 38 | if byt[i] != []byte(tt.str)[i] { 39 | t.Errorf("expected ih.Bytes() to equal %s, got %s", []byte(tt.str), ih.Bytes()) 40 | } 41 | } 42 | } else { 43 | if err == nil { 44 | t.Errorf("FromString should have failed for %s", tt.str) 45 | } 46 | } 47 | } 48 | } 49 | 50 | func TestInfohashLength(t *testing.T) { 51 | ih := GenInfohash() 52 | if len(ih) != 20 { 53 | t.Errorf("%s as string should be length 20, got %d", ih, len(ih)) 54 | } 55 | } 56 | 57 | func TestInfohashDistance(t *testing.T) { 58 | id := "d1c5676ae7ac98e8b19f63565905105e3c4c37a2" 59 | 60 | var tests = []struct { 61 | ih string 62 | other string 63 | distance int 64 | }{ 65 | {id, id, 160}, 66 | {id, "d1c5676ae7ac98e8b19f63565905105e3c4c37a3", 159}, 67 | } 68 | 69 | ih, err := InfohashFromString(id) 70 | if err != nil { 71 | t.Errorf("Failed to create Infohash: %s", err) 72 | } 73 | 74 | for _, tt := range tests { 75 | other, err := InfohashFromString(tt.other) 76 | if err != nil { 77 | t.Errorf("Failed to create Infohash: %s", err) 78 | } 79 | 80 | dist := ih.Distance(*other) 81 | if dist != tt.distance { 82 | t.Errorf("Distance() => %d, expected %d", dist, tt.distance) 83 | } 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /models/peer.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import ( 4 | "fmt" 5 | "net" 6 | "time" 7 | ) 8 | 9 | // Peer on DHT network 10 | type Peer struct { 11 | Addr net.Addr `db:"address"` 12 | Infohash Infohash `db:"infohash"` 13 | Created time.Time `db:"created" json:"created"` 14 | Updated time.Time `db:"updated" json:"updated"` 15 | } 16 | 17 | // String implements fmt.Stringer 18 | func (p Peer) String() string { 19 | return fmt.Sprintf("%s (%s)", p.Infohash, p.Addr.String()) 20 | } 21 | -------------------------------------------------------------------------------- /models/storage.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import () 4 | 5 | type migratable interface { 6 | MigrateSchema() error 7 | } 8 | 9 | type torrentSearcher interface { 10 | TorrentsByHash(hash Infohash) (*Torrent, error) 11 | TorrentsByName(query string, offset, limit int) ([]*Torrent, error) 12 | TorrentsByTags(tags []string, offset, limit int) ([]*Torrent, error) 13 | } 14 | 15 | type PeerStore interface { 16 | SavePeer(*Peer) error 17 | RemovePeer(*Peer) error 18 | } 19 | 20 | type TorrentStore interface { 21 | SaveTorrent(*Torrent) error 22 | RemoveTorrent(*Torrent) error 23 | RemovePeer(*Peer) error 24 | } 25 | 26 | type InfohashStore interface { 27 | PendingInfohashes(int) ([]*Peer, error) 28 | } 29 | -------------------------------------------------------------------------------- /models/tag.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | type tagStore interface { 4 | saveTag(string) (int, error) 5 | } 6 | -------------------------------------------------------------------------------- /models/torrent.go: -------------------------------------------------------------------------------- 1 | package models 2 | 3 | import ( 4 | "bytes" 5 | "crypto/sha1" 6 | "fmt" 7 | "os" 8 | "strings" 9 | "time" 10 | 11 | "src.userspace.com.au/dhtsearch/krpc" 12 | "src.userspace.com.au/go-bencode" 13 | ) 14 | 15 | // Data for persistent storage 16 | type Torrent struct { 17 | ID int `json:"-"` 18 | Infohash Infohash `json:"infohash"` 19 | Name string `json:"name"` 20 | Files []File `json:"files" db:"-"` 21 | Size int `json:"size"` 22 | Updated time.Time `json:"updated"` 23 | Created time.Time `json:"created"` 24 | Tags []string `json:"tags" db:"-"` 25 | } 26 | 27 | type File struct { 28 | ID int `json:"-"` 29 | Path string `json:"path"` 30 | Size int `json:"size"` 31 | TorrentID int `json:"torrent_id" db:"torrent_id"` 32 | } 33 | 34 | func InfohashMatchesMetadata(ih Infohash, md []byte) bool { 35 | info := sha1.Sum(md) 36 | return bytes.Equal([]byte(ih), info[:]) 37 | } 38 | 39 | func TorrentFromMetadata(ih Infohash, md []byte) (*Torrent, error) { 40 | if !InfohashMatchesMetadata(ih, md) { 41 | return nil, fmt.Errorf("infohash does not match metadata") 42 | } 43 | info, _, err := bencode.DecodeDict(md, 0) 44 | if err != nil { 45 | return nil, err 46 | } 47 | 48 | // Get the directory or advisory filename 49 | name, err := krpc.GetString(info, "name") 50 | if err != nil { 51 | return nil, err 52 | } 53 | 54 | bt := Torrent{ 55 | Infohash: ih, 56 | Name: name, 57 | } 58 | 59 | if files, err := krpc.GetList(info, "files"); err == nil { 60 | // Multiple file mode 61 | bt.Files = make([]File, len(files)) 62 | 63 | // Files is a list of dicts 64 | for i, item := range files { 65 | file := item.(map[string]interface{}) 66 | 67 | // Paths is a list of strings 68 | paths := file["path"].([]interface{}) 69 | path := make([]string, len(paths)) 70 | for j, p := range paths { 71 | path[j] = p.(string) 72 | } 73 | 74 | fSize, err := krpc.GetInt(file, "length") 75 | if err != nil { 76 | return nil, err 77 | } 78 | bt.Files[i] = File{ 79 | // Assume Unix path sep? 80 | Path: strings.Join(path[:], string(os.PathSeparator)), 81 | Size: fSize, 82 | } 83 | // Ensure the torrent size totals all files' 84 | bt.Size = bt.Size + fSize 85 | } 86 | } else if length, err := krpc.GetInt(info, "length"); err == nil { 87 | // Single file mode 88 | bt.Size = length 89 | } else { 90 | return nil, fmt.Errorf("found neither length or files") 91 | } 92 | return &bt, nil 93 | } 94 | --------------------------------------------------------------------------------