diff options
Diffstat (limited to 'src/Network/BitTorrent/Exchange/Connection.hs')
-rw-r--r-- | src/Network/BitTorrent/Exchange/Connection.hs | 111 |
1 files changed, 105 insertions, 6 deletions
diff --git a/src/Network/BitTorrent/Exchange/Connection.hs b/src/Network/BitTorrent/Exchange/Connection.hs index dde9a468..fd9022da 100644 --- a/src/Network/BitTorrent/Exchange/Connection.hs +++ b/src/Network/BitTorrent/Exchange/Connection.hs | |||
@@ -49,6 +49,19 @@ module Network.BitTorrent.Exchange.Connection | |||
49 | , connSession | 49 | , connSession |
50 | , connStats | 50 | , connStats |
51 | 51 | ||
52 | -- ** Status | ||
53 | , PeerStatus (..) | ||
54 | , ConnectionStatus (..) | ||
55 | , updateStatus | ||
56 | , statusUpdates | ||
57 | , clientStatus | ||
58 | , remoteStatus | ||
59 | , canUpload | ||
60 | , canDownload | ||
61 | , defaultUnchokeSlots | ||
62 | , defaultRechokeInterval | ||
63 | |||
64 | |||
52 | -- * Setup | 65 | -- * Setup |
53 | , ConnectionPrefs (..) | 66 | , ConnectionPrefs (..) |
54 | , SessionLink (..) | 67 | , SessionLink (..) |
@@ -102,21 +115,21 @@ import Control.Monad.State | |||
102 | import Control.Lens | 115 | import Control.Lens |
103 | import Data.ByteString as BS | 116 | import Data.ByteString as BS |
104 | import Data.ByteString.Lazy as BSL | 117 | import Data.ByteString.Lazy as BSL |
105 | import Data.Conduit | 118 | import Data.Conduit as C |
106 | import Data.Conduit.Cereal | 119 | import Data.Conduit.Cereal |
107 | import Data.Conduit.List | 120 | import Data.Conduit.List |
108 | import Data.Conduit.Network | 121 | import Data.Conduit.Network |
109 | import Data.Default | 122 | import Data.Default |
110 | import Data.IORef | 123 | import Data.IORef |
111 | import Data.List as L | 124 | import Data.List as L |
112 | import Data.Maybe | 125 | import Data.Maybe as M |
113 | import Data.Monoid | 126 | import Data.Monoid |
114 | import Data.Serialize as S | 127 | import Data.Serialize as S |
115 | import Data.Typeable | 128 | import Data.Typeable |
116 | import Network | 129 | import Network |
117 | import Network.Socket hiding (Connected) | 130 | import Network.Socket hiding (Connected) |
118 | import Network.Socket.ByteString as BS | 131 | import Network.Socket.ByteString as BS |
119 | import Text.PrettyPrint as PP hiding (($$), (<>)) | 132 | import Text.PrettyPrint as PP hiding ((<>)) |
120 | import Text.PrettyPrint.Class | 133 | import Text.PrettyPrint.Class |
121 | import Text.Show.Functions () | 134 | import Text.Show.Functions () |
122 | import System.Log.FastLogger (ToLogStr(..)) | 135 | import System.Log.FastLogger (ToLogStr(..)) |
@@ -125,7 +138,6 @@ import System.Timeout | |||
125 | import Data.Torrent.Bitfield as BF | 138 | import Data.Torrent.Bitfield as BF |
126 | import Data.Torrent.InfoHash | 139 | import Data.Torrent.InfoHash |
127 | import Network.BitTorrent.Core | 140 | import Network.BitTorrent.Core |
128 | import Network.BitTorrent.Exchange.Connection.Status | ||
129 | import Network.BitTorrent.Exchange.Message as Msg | 141 | import Network.BitTorrent.Exchange.Message as Msg |
130 | 142 | ||
131 | -- TODO handle port message? | 143 | -- TODO handle port message? |
@@ -463,6 +475,93 @@ instance Default Options where | |||
463 | } | 475 | } |
464 | 476 | ||
465 | {----------------------------------------------------------------------- | 477 | {----------------------------------------------------------------------- |
478 | -- Peer status | ||
479 | -----------------------------------------------------------------------} | ||
480 | |||
481 | -- | Connections contain two bits of state on either end: choked or | ||
482 | -- not, and interested or not. | ||
483 | data PeerStatus = PeerStatus | ||
484 | { -- | Choking is a notification that no data will be sent until | ||
485 | -- unchoking happens. | ||
486 | _choking :: !Bool | ||
487 | |||
488 | -- | | ||
489 | , _interested :: !Bool | ||
490 | } deriving (Show, Eq, Ord) | ||
491 | |||
492 | $(makeLenses ''PeerStatus) | ||
493 | |||
494 | instance Pretty PeerStatus where | ||
495 | pretty PeerStatus {..} = | ||
496 | pretty (Choking _choking) <+> "and" <+> pretty (Interested _interested) | ||
497 | |||
498 | -- | Connections start out choked and not interested. | ||
499 | instance Default PeerStatus where | ||
500 | def = PeerStatus True False | ||
501 | |||
502 | instance Monoid PeerStatus where | ||
503 | mempty = def | ||
504 | mappend a b = PeerStatus | ||
505 | { _choking = _choking a && _choking b | ||
506 | , _interested = _interested a || _interested b | ||
507 | } | ||
508 | |||
509 | -- | Can be used to update remote peer status using incoming 'Status' | ||
510 | -- message. | ||
511 | updateStatus :: StatusUpdate -> PeerStatus -> PeerStatus | ||
512 | updateStatus (Choking b) = choking .~ b | ||
513 | updateStatus (Interested b) = interested .~ b | ||
514 | |||
515 | -- | Can be used to generate outcoming messages. | ||
516 | statusUpdates :: PeerStatus -> PeerStatus -> [StatusUpdate] | ||
517 | statusUpdates a b = M.catMaybes $ | ||
518 | [ if _choking a == _choking b then Nothing | ||
519 | else Just $ Choking $ _choking b | ||
520 | , if _interested a == _interested b then Nothing | ||
521 | else Just $ Interested $ _interested b | ||
522 | ] | ||
523 | |||
524 | {----------------------------------------------------------------------- | ||
525 | -- Connection status | ||
526 | -----------------------------------------------------------------------} | ||
527 | |||
528 | -- | Status of the both endpoints. | ||
529 | data ConnectionStatus = ConnectionStatus | ||
530 | { _clientStatus :: !PeerStatus | ||
531 | , _remoteStatus :: !PeerStatus | ||
532 | } deriving (Show, Eq) | ||
533 | |||
534 | $(makeLenses ''ConnectionStatus) | ||
535 | |||
536 | instance Pretty ConnectionStatus where | ||
537 | pretty ConnectionStatus {..} = | ||
538 | "this " PP.<+> pretty _clientStatus PP.$$ | ||
539 | "remote" PP.<+> pretty _remoteStatus | ||
540 | |||
541 | -- | Connections start out choked and not interested. | ||
542 | instance Default ConnectionStatus where | ||
543 | def = ConnectionStatus def def | ||
544 | |||
545 | -- | Can the client transfer to the remote peer? | ||
546 | canUpload :: ConnectionStatus -> Bool | ||
547 | canUpload ConnectionStatus {..} | ||
548 | = _interested _remoteStatus && not (_choking _clientStatus) | ||
549 | |||
550 | -- | Can the client transfer from the remote peer? | ||
551 | canDownload :: ConnectionStatus -> Bool | ||
552 | canDownload ConnectionStatus {..} | ||
553 | = _interested _clientStatus && not (_choking _remoteStatus) | ||
554 | |||
555 | -- | Indicates how many peers are allowed to download from the client | ||
556 | -- by default. | ||
557 | defaultUnchokeSlots :: Int | ||
558 | defaultUnchokeSlots = 4 | ||
559 | |||
560 | -- | | ||
561 | defaultRechokeInterval :: Int | ||
562 | defaultRechokeInterval = 10 * 1000 * 1000 | ||
563 | |||
564 | {----------------------------------------------------------------------- | ||
466 | -- Connection | 565 | -- Connection |
467 | -----------------------------------------------------------------------} | 566 | -----------------------------------------------------------------------} |
468 | 567 | ||
@@ -681,7 +780,7 @@ runWire action sock chan conn = flip runReaderT conn $ runConnected $ | |||
681 | conduitGet S.get $= | 780 | conduitGet S.get $= |
682 | trackFlow RemotePeer $= | 781 | trackFlow RemotePeer $= |
683 | action $= | 782 | action $= |
684 | trackFlow ThisPeer $$ | 783 | trackFlow ThisPeer C.$$ |
685 | sinkChan chan | 784 | sinkChan chan |
686 | 785 | ||
687 | -- | This function will block until a peer send new message. You can | 786 | -- | This function will block until a peer send new message. You can |
@@ -835,7 +934,7 @@ closePending PendingConnection {..} = do | |||
835 | 934 | ||
836 | chanToSock :: Int -> Chan Message -> Socket -> IO () | 935 | chanToSock :: Int -> Chan Message -> Socket -> IO () |
837 | chanToSock ka chan sock = | 936 | chanToSock ka chan sock = |
838 | sourceChan ka chan $= conduitPut S.put $$ sinkSocket sock | 937 | sourceChan ka chan $= conduitPut S.put C.$$ sinkSocket sock |
839 | 938 | ||
840 | afterHandshaking :: ChannelSide -> PeerAddr IP -> Socket -> HandshakePair | 939 | afterHandshaking :: ChannelSide -> PeerAddr IP -> Socket -> HandshakePair |
841 | -> ConnectionConfig s -> IO () | 940 | -> ConnectionConfig s -> IO () |