Ton

1
{-
2
  Originally created by:
3
  /------------------------------------------------------------------------\
4
  | Created for: Telegram (Open Network) Blockchain Contest                |
5
  |      Task 3: DNS Resolver (Manually controlled)                        |
6
  >------------------------------------------------------------------------<
7
  | Author: Oleksandr Murzin (tg: @skydev / em: alexhacker64@gmail.com)    |
8
  |         October 2019                                                   |
9
  \------------------------------------------------------------------------/
10
  Updated to actual DNS standard version by starlightduck in 2022
11
-}
12

13
;;===========================================================================;;
14
;; Custom ASM instructions                                                   ;;
15
;;===========================================================================;;
16

17
cell udict_get_ref_(cell dict, int key_len, int index) asm(index dict key_len) "DICTUGETOPTREF";
18

19
(cell, ()) pfxdict_set_ref(cell dict, int key_len, slice key, cell value) {
20
  throw_unless(33, dict~pfxdict_set?(key_len, key, begin_cell().store_maybe_ref(value).end_cell().begin_parse()));
21
  return (dict, ());
22
}
23

24
(slice, cell, slice, int) pfxdict_get_ref(cell dict, int key_len, slice key) inline_ref {
25
  (slice pfx, slice val, slice tail, int succ) = dict.pfxdict_get?(key_len, key);
26
  cell res = succ ? val~load_maybe_ref() : null();
27
  return (pfx, res, tail, succ);
28
}
29

30
;;===========================================================================;;
31
;; Utility functions                                                         ;;
32
;;===========================================================================;;
33

34
(int, int, int, cell, cell) load_data() inline_ref {
35
  slice cs = get_data().begin_parse();
36
  var res = (cs~load_uint(32), cs~load_uint(64), cs~load_uint(256), cs~load_dict(), cs~load_dict());
37
  cs.end_parse();
38
  return res;
39
}
40

41
() store_data(int contract_id, int last_cleaned, int public_key, cell root, old_queries) impure {
42
  set_data(begin_cell()
43
      .store_uint(contract_id, 32)
44
      .store_uint(last_cleaned, 64)
45
      .store_uint(public_key, 256)
46
      .store_dict(root)
47
      .store_dict(old_queries)
48
      .end_cell());
49
}
50

51
;;===========================================================================;;
52
;; Internal message handler (Code 0)                                         ;;
53
;;===========================================================================;;
54

55
() recv_internal(int msg_value, cell in_msg_cell, slice in_msg) impure {
56
  ;; not interested at all
57
}
58

59
;;===========================================================================;;
60
;; External message handler (Code -1)                                        ;;
61
;;===========================================================================;;
62

63
{-
64
  External message structure:
65
    [Bytes<512b>:signature] [UInt<32b>:seqno] [UInt<6b>:operation]
66
    [Either b0: inline name (<= 58-x Bytes) or b1: reference-stored name)
67
                                   x depends on operation
68
    Use of 6-bit op instead of 32-bit allows to save 4 bytes for inline name
69
  Inline [Name] structure: [UInt<6b>:length] [Bytes<lengthB>:data]
70
  Operations (continuation of message):
71
  00 Contract initialization message (only if seqno = 0) (x=-)
72
  11 VSet: set specified value to specified subdomain->category (x=2)
73
    [UInt<256b>:category] [Name<?>:subdomain] [Cell<1r>:value]
74
  12 VDel: delete specified subdomain->category (x=2)
75
    [UInt<256b>:category] [Name<?>:subdomain]
76
  21 DSet: replace entire category dictionary of domain with provided (x=0)
77
    [Name<?>:subdomain] [Cell<1r>:new_cat_table]
78
  22 DDel: delete entire category dictionary of specified domain (x=0)
79
    [Name<?>:subdomain]
80
  31 TSet: replace ENTIRE DOMAIN TABLE with the provided tree root cell (x=-)
81
    [Cell<1r>:new_domains_table]
82
  32 TDel: nullify ENTIRE DOMAIN TABLE (x=-)
83
  51 OSet: replace owner public key with a new one (x=-)
84
    [UInt<256b>:new_public_key]
85
-}
86

87
() after_code_upgrade(cell root, slice ops, cont old_code) impure method_id(1666);
88

89
(cell, slice) process_op(cell root, slice ops) inline_ref {
90
  int op = ops~load_uint(6);
91
  if (op < 10) {
92
    ifnot (op) {
93
      ;; 00 Noop: No operation
94
      return (root, ops);
95
    }
96
    if (op == 1) {
97
      ;; 01 SMsg: Send Message
98
      var mode = ops~load_uint(8);
99
      send_raw_message(ops~load_ref(), mode);
100
      return (root, ops);
101
    }
102
    if (op == 9) {
103
      ;; 09 CodeUpgrade
104
      var new_code = ops~load_ref();
105
      set_code(new_code);
106
      var old_code = get_c3();
107
      set_c3(new_code.begin_parse().bless());
108
      after_code_upgrade(root, ops, old_code);
109
      throw(0);
110
      return (root, ops);
111
    }
112
    throw(45);
113
    return (root, ops);
114
  }
115
  int cat = 0;
116
  if (op < 20) {
117
    ;; for operations with codes 10..19 category is required
118
    cat = ops~load_uint(256); ;; update: category length now u256 instead of i16
119
  }
120
  slice name = null();   ;; any slice value
121
  cell cat_table = null();
122
  if (op < 30) {
123
    ;; for operations with codes 10..29 name is required
124
    int is_name_ref = (ops~load_uint(1) == 1);
125
    if (is_name_ref) {
126
      ;; name is stored in separate referenced cell
127
      name = ops~load_ref().begin_parse();
128
    } else {
129
      ;; name is stored inline
130
      int name_len = ops~load_uint(6) * 8;
131
      name = ops~load_bits(name_len);
132
    }
133
    ;; at least one character not counting \0
134
    throw_unless(38, name.slice_bits() >= 16);
135
    ;; name shall end with \0
136
    int name_last_byte = name.slice_last(8).preload_uint(8);
137
    throw_if(40, name_last_byte);
138
    ;; count zero separators
139
    int zeros = 0;
140
    slice cname = name;
141
    repeat (cname.slice_bits() ^>> 3) {
142
      int c = cname~load_uint(8);
143
      zeros -= (c == 0);
144
    }
145
    ;; throw_unless(39, zeros == 1);
146
    name = begin_cell().store_uint(zeros, 7).store_slice(name).end_cell().begin_parse();
147
  }
148
  ;; operation with codes 10..19 manipulate category dict
149
  ;; lets try to find it and store into a variable
150
  ;; operations with codes 20..29 replace / delete dict, no need
151
  if (op < 20) {
152
    ;; lets resolve the name here so as not to duplicate the code
153
    (slice pfx, cell val, slice tail, int succ) = 
154
      root.pfxdict_get_ref(1023, name);
155
    if (succ) {
156
      ;; must match EXACTLY to prevent accident changes
157
      throw_unless(35, tail.slice_empty?()); 
158
      cat_table = val;
159
    }
160
    ;; otherwise cat_table is null which is reasonable for actions
161
  }
162
  ;; 11 VSet: set specified value to specified subdomain->category
163
  if (op == 11) {
164
    cell new_value = ops~load_maybe_ref();
165
    cat_table~udict_set_get_ref(256, cat, new_value); ;; update: category length now u256 instead of i16
166
    root~pfxdict_set_ref(1023, name, cat_table);
167
    return (root, ops);
168
  }
169
  ;; 12 VDel: delete specified subdomain->category value
170
  if (op == 12) {
171
    if (cat_table~udict_delete?(256, cat)) { ;; update: category length now u256 instead of i16
172
       root~pfxdict_set_ref(1023, name, cat_table);      
173
    }
174
    return (root, ops);
175
  }
176
  ;; 21 DSet: replace entire category dictionary of domain with provided
177
  if (op == 21) {
178
    cell new_cat_table = ops~load_maybe_ref();
179
    root~pfxdict_set_ref(1023, name, new_cat_table);
180
    return (root, ops);
181
  }
182
  ;; 22 DDel: delete entire category dictionary of specified domain
183
  if (op == 22) {
184
    root~pfxdict_delete?(1023, name);
185
    return (root, ops);
186
  }
187
  ;; 31 TSet: replace ENTIRE DOMAIN TABLE with the provided tree root cell
188
  if (op == 31) {
189
    cell new_tree_root = ops~load_maybe_ref();
190
    ;; no sanity checks cause they would cost immense gas
191
    return (new_tree_root, ops);
192
  }
193
  ;; 32 TDel: nullify ENTIRE DOMAIN TABLE
194
  if (op == 32) {
195
    return (null(), ops);
196
  }
197
  throw(44); ;; invalid operation
198
  return (null(), ops);
199
}
200

201
cell process_ops(cell root, slice ops) inline_ref {
202
  var stop = false;
203
  root~touch();
204
  ops~touch();
205
  do {
206
    (root, ops) = process_op(root, ops);
207
    if (ops.slice_data_empty?()) {
208
      if (ops.slice_refs()) {
209
        ops = ops~load_ref().begin_parse();
210
      } else {
211
        stop = true;
212
      }
213
    }
214
  } until (stop);
215
  return root;
216
}
217

218
() recv_external(slice in_msg) impure {
219
  ;; Load data
220
  (int contract_id, int last_cleaned, int public_key, cell root, cell old_queries) = load_data();
221

222
  ;; validate signature and seqno
223
  slice signature = in_msg~load_bits(512);
224
  int shash = slice_hash(in_msg);
225
  var (query_contract, query_id) = (in_msg~load_uint(32), in_msg~load_uint(64));
226
  var bound = (now() << 32);
227
  throw_if(35, query_id < bound);
228
  (_, var found?) = old_queries.udict_get?(64, query_id);
229
  throw_if(32, found?);
230
  throw_unless(34, contract_id == query_contract);
231
  throw_unless(35, check_signature(shash, signature, public_key));
232
  accept_message(); ;; message is signed by owner, sanity not guaranteed yet
233

234
  int op = in_msg.preload_uint(6);
235
  if (op == 51) {
236
    in_msg~skip_bits(6);
237
    public_key = in_msg~load_uint(256);
238
  } else {
239
    root = process_ops(root, in_msg);
240
  }
241

242
  bound -= (64 << 32);   ;; clean up records expired more than 64 seconds ago
243
  old_queries~udict_set_builder(64, query_id, begin_cell());
244
  var queries = old_queries;
245
  do {
246
    var (old_queries', i, _, f) = old_queries.udict_delete_get_min(64);
247
    f~touch();
248
    if (f) {
249
      f = (i < bound);
250
    }
251
    if (f) {
252
      old_queries = old_queries';
253
      last_cleaned = i;
254
    }
255
  } until (~ f);
256

257
  store_data(contract_id, last_cleaned, public_key, root, old_queries);
258
}
259

260
() after_code_upgrade(cell root, slice ops, cont old_code) impure method_id(1666) {
261
}
262

263
{-
264
  Data structure:
265
  Root cell: [UInt<32b>:seqno] [UInt<256b>:owner_public_key] 
266
         [OptRef<1b+1r?>:Hashmap<PfxDict:Slice->CatTable>:domains]
267
  <CatTable> := HashmapE 256 (~~16~~) ^DNSRecord
268
  
269
  STORED DOMAIN NAME SLICE FORMAT: (#ZeroChars<7b>) (Domain name value)
270
  #Zeros allows to simultaneously store, for example, com\0 and com\0google\0
271
  That will be stored as \1com\0 and \2com\0google\0 (pfx tree has restricitons)
272
  This will allow to resolve more specific requests to subdomains, and resort
273
      to parent domain next resolver lookup if subdomain is not found
274
  com\0goo\0 lookup will, for example look up \2com\0goo\0 and then
275
      \1com\0goo\0 which will return \1com\0 (as per pfx tree) with -1 cat
276
-}
277

278
;;===========================================================================;;
279
;; Getter methods                                                            ;;
280
;;===========================================================================;;
281

282
;; Retrieve contract id (in case several contracts are managed with the same private key)
283
int get_contract_id() method_id {
284
  return get_data().begin_parse().preload_uint(32);
285
}
286

287
int get_public_key() method_id {
288
  var cs = get_data().begin_parse();
289
  cs~load_uint(32 + 64);
290
  return cs.preload_uint(256);
291
}
292

293
;;8m  dns-record-value
294
(int, cell) dnsresolve(slice subdomain, int category) method_id {
295
  int bits = subdomain.slice_bits();
296
  ifnot (bits) {
297
    ;; return (0, null());  ;; zero-length input
298
    throw(30); ;; update: throw exception for empty input
299
  }
300
  throw_if(30, bits & 7); ;; malformed input (~ 8n-bit)
301

302
  int name_last_byte = subdomain.slice_last(8).preload_uint(8);
303
  if (name_last_byte) {
304
    subdomain = begin_cell().store_slice(subdomain) ;; append zero byte
305
                            .store_uint(0, 8).end_cell().begin_parse();
306
    bits += 8;
307
  }
308
  if (bits == 8) {
309
    return (8, null()); ;; zero-length input, but with zero byte
310
    ;; update: return 8 as resolved, but with no data
311
  }
312
  int name_first_byte = subdomain.preload_uint(8);
313
  if (name_first_byte == 0) {
314
    ;; update: remove prefix \0
315
    subdomain~load_uint(8);
316
    bits -= 8;
317
  }
318
  (_, _, _, cell root, _) = load_data();
319
  
320
  slice cname = subdomain;
321
  int zeros = 0;
322
  repeat (bits >> 3) {
323
    int c = cname~load_uint(8);
324
    zeros -= (c == 0);
325
  }
326

327
  ;; can't move these declarations lower, will cause errors!
328
  slice pfx = cname;
329
  cell val = null();
330
  slice tail = cname;
331

332
  do {
333
    slice pfxname = begin_cell().store_uint(zeros, 7)
334
      .store_slice(subdomain).end_cell().begin_parse();
335
    (pfx, val, tail, int succ) = root.pfxdict_get_ref(1023, pfxname);
336
    zeros = succ ^ (zeros - 1);   ;; break on success
337
  } until (zeros <= 0);
338

339
  ifnot (zeros) {
340
    return (0, null()); ;; failed to find entry in prefix dictionary
341
  }
342

343
  zeros = - zeros;
344

345
  ifnot (tail.slice_empty?()) { ;; if we have tail then len(pfx) < len(subdomain)
346
    ;; incomplete subdomain found, must return next resolver
347
    category = "dns_next_resolver"H; ;; 0x19f02441ee588fdb26ee24b2568dd035c3c9206e11ab979be62e55558a1d17ff
348
    ;; update: next resolver is now sha256("dns_next_resolver") instead of -1
349
  }
350
  int pfx_bits = pfx.slice_bits() - 7;
351
  cell cat_table = val;
352
  ;; pfx.slice_bits() will contain 8m, where m is number of bytes in subdomain 
353
  ;;   COUNTING the zero byte (if structurally correct: no multiple-ZB keys)
354
  ;;   which corresponds to 8m, m=one plus the number of bytes in the subdomain found)
355
  if (category == 0) {
356
    return (pfx_bits, cat_table); ;; return cell with entire dictionary for 0
357
  } else {
358
    cell cat_found = cat_table.udict_get_ref_(256, category); ;; update: category length now u256 instead of i16
359
    return (pfx_bits, cat_found);
360
  }
361
}
362