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
|