aboutsummaryrefslogtreecommitdiff
path: root/lib/client/shardManager.ml
blob: cf849bcd09ef9690cb592b816cb647a62baf78d8 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
open Lwt.Infix
open Websocket

module Opcode = struct
    type t =
        | DISPATCH
        | HEARTBEAT
        | IDENTIFY
        | STATUS_UPDATE
        | VOICE_STATE_UPDATE
        | RESUME
        | RECONNECT
        | REQUEST_GUILD_MEMBERS
        | INVALID_SESSION
        | HELLO
        | HEARTBEAT_ACK

    let to_int = function
        | DISPATCH -> 0
        | HEARTBEAT -> 1
        | IDENTIFY -> 2
        | STATUS_UPDATE -> 3
        | VOICE_STATE_UPDATE -> 4
        | RESUME -> 6
        | RECONNECT -> 7
        | REQUEST_GUILD_MEMBERS -> 8
        | INVALID_SESSION -> 9
        | HELLO -> 10
        | HEARTBEAT_ACK -> 11

    let to_string = function
        | DISPATCH -> "DISPATCH"
        | HEARTBEAT -> "HEARTBEAT"
        | IDENTIFY -> "IDENTIFY"
        | STATUS_UPDATE -> "STATUS_UPDATE"
        | VOICE_STATE_UPDATE -> "VOICE_STATE_UPDATE"
        | RESUME -> "RESUME"
        | RECONNECT -> "RECONNECT"
        | REQUEST_GUILD_MEMBERS -> "REQUEST_GUILD_MEMBER"
        | INVALID_SESSION -> "INVALID_SESSION"
        | HELLO -> "HELLO"
        | HEARTBEAT_ACK -> "HEARTBEAT_ACK"
end

module ShardMap = Map.Make(
    struct
        type t = int
        let compare: int -> int -> int = Pervasives.compare
    end
)

module Shard = struct
    type t = {
        send: (Frame.t -> unit Lwt.t);
        recv: (unit -> Frame.t Lwt. t);
        id: int;
        hb_interval: int;
        session_id: string;
        seq: int;
    }

    let send payload shard =
        payload
        |> shard.send

    let wrap_payload d op =
        `Assoc [
            ("op", `Int op);
            ("d", d)
        ]

    let create_frame content =
        Frame.create ~content ()

    let identify ?(threshold=250) shard total token =
        let p = wrap_payload (`Assoc [
            ("token", `String token);
            ("properties", `Assoc [
                ("$os", `String Sys.os_type);
                ("$browser", `String "animus");
                ("$device", `String "animus");
            ]);
            ("large_threshold", `Int threshold);
            ("shard", `List [`Int shard.id; `Int total]);
        ]) (Opcode.to_int IDENTIFY) in
        let p = create_frame (Yojson.Basic.to_string p) in
        send p shard

    let resume shard token =
        let p = wrap_payload (`Assoc [
            ("token", `String token);
            ("session_id", `String shard.session_id);
            ("seq", `Int shard.seq);
        ]) (Opcode.to_int RESUME) in
        let p = create_frame (Yojson.Basic.to_string p) in
        send p shard

    let heartbeat shard =
        let p = wrap_payload (`Int shard.seq) (Opcode.to_int HEARTBEAT) in
        let p = create_frame (Yojson.Basic.to_string p) in
        send p shard

    (* Use options *)
    let connect ~options uri id total token =
        let url = Uri.to_string uri in
        let ip = Ipaddr.V4 Ipaddr.V4.any in
        Websocket_lwt.with_connection (`TLS (`Hostname url, `IP ip, `Port 443)) uri
        >|= fun (recv, send) ->
        let shard = { send; recv; hb_interval = 42500; id; session_id = ""; seq = 0; } in
        heartbeat shard >>= (fun _ ->
            identify shard total token (* This needs to be handled, I'm just pleasing the compiler *)
        ) |> ignore;
        shard
end

type t = {
    shards: Shard.t ShardMap.t;
    gateway_url: Uri.t;
    token: string;
}

let create_shard ?(options=[]) manager =
    let id = (ShardMap.cardinal manager.shards) + 1 in
    Shard.connect ~options manager.gateway_url id (ShardMap.cardinal manager.shards) manager.token
    >|= fun shard ->
    ShardMap.add id shard manager.shards

let update_shard id shard manager =
    match ShardMap.exists (fun k _ -> id == k) manager.shards with
    | true -> ShardMap.add id shard manager.shards
    | false -> manager.shards

let heartbeat manager id =
    let shard = ShardMap.find id manager.shards in
    Shard.heartbeat shard

let identify manager id =
    let total = ShardMap.cardinal manager.shards in
    let shard = ShardMap.find id manager.shards in
    Shard.identify shard total manager.token

let resume manager id =
    let shard = ShardMap.find id manager.shards in
    Shard.resume shard manager.token