1 | /* Copyright 2022 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | // This file implements random_index_shuffle() by using a simple block chiper |
17 | // for pseudorandom permutations. |
18 | // This idea is described as cycle-walking cipher in |
19 | // https://www.cs.ucdavis.edu/~rogaway/papers/subset.pdf |
20 | // |
21 | // We use the Simon block cipher described in |
22 | // https://eprint.iacr.org/2013/404 |
23 | // and following recommendations in |
24 | // https://nsacyber.github.io/simon-speck/implementations/ImplementationGuide1.1.pdf. |
25 | // However we use a single fixed key size and support arbtitary block sizes. |
26 | // Further we fixed the number of rounds in the Feistel structuro to be always |
27 | // 4. This reduces the computational cost and still gives good shuffle behavior. |
28 | // |
29 | // Warning: Given the modifications descripted above this implementation should |
30 | // not be used for application that require cryptograhic secure RNGs. |
31 | |
32 | #include "tensorflow/core/kernels/random_index_shuffle.h" |
33 | |
34 | #include <assert.h> |
35 | |
36 | #include <algorithm> |
37 | #include <array> |
38 | #include <bitset> |
39 | #include <cmath> |
40 | |
41 | #include "tensorflow/core/platform/types.h" |
42 | |
43 | namespace tensorflow { |
44 | namespace random { |
45 | // Some of the macros below require a minimum word size of 8 |
46 | // (2 word size = block size). |
47 | // Smaller block sizes might give poor results in terms of randomness. |
48 | constexpr int kMinBlockSize = 16; |
49 | |
50 | namespace impl { |
51 | |
52 | #define ROTL(x, r, W) (((x) << (r)) | (x >> (W - (r)))) |
53 | #define ROTR(x, r, W) (((x) >> (r)) | ((x) << (W - (r)))) |
54 | #define SIMON_F(x, W) ((ROTL(x, 1, W) & ROTL(x, 8, W) ^ ROTL(x, 2, W))) |
55 | #define SIMON_Rx2(x, y, k1, k2, W) \ |
56 | (y ^= SIMON_F(x, W), y ^= k1, x ^= SIMON_F(y, W), x ^= k2) |
57 | |
58 | // Returns the keys per round for a Simon cipher. |
59 | // This variant uses std::bitset and can generate keys with any number of bits. |
60 | // This should not be used to encrypt data. It's not secure. We only use it |
61 | // to generate pseudorandom permutations. |
62 | template <int W> |
63 | std::vector<std::bitset<W>> simon_key_schedule( |
64 | const std::array<uint32_t, 3>& key, const int32_t rounds) { |
65 | // Required by ROTR/ROTL |
66 | static_assert(W >= 8, "Minimum word size is 8 bits." ); |
67 | const auto c = std::bitset<W>(0xfffffffc); |
68 | auto z = std::bitset<W>(0x7369f885192c0ef5LL); |
69 | std::vector<std::bitset<W>> rk({key[0], key[1], key[2]}); |
70 | rk.reserve(rounds); |
71 | for (int i = 3; i < rounds; i++) { |
72 | rk.push_back(c ^ (z & std::bitset<W>(1)) ^ rk[i - 3] ^ |
73 | ROTR(rk[i - 1], 3, W) ^ ROTR(rk[i - 1], 4, W)); |
74 | z >>= 1; |
75 | } |
76 | return rk; |
77 | } |
78 | |
79 | // Encrypts the given value using the Simon chipher. |
80 | // This is should not be used to encrypt data. It's not secure. We only use it |
81 | // to generate pseudorandom permutations. |
82 | template <int W> |
83 | uint64_t simon_encrypt(const uint64_t value, |
84 | const std::vector<std::bitset<W>>& round_keys) { |
85 | // Required by ROTR/ROTL |
86 | static_assert(W >= 8, "Minimum word size is 8 bits." ); |
87 | std::bitset<W> left(value >> W); |
88 | std::bitset<W> right(value); |
89 | for (int i = 0; i < round_keys.size();) { |
90 | SIMON_Rx2(right, left, round_keys[i++], round_keys[i++], W); |
91 | } |
92 | return (left.to_ullong() << W) | right.to_ullong(); |
93 | } |
94 | |
95 | // In the original implementation the number of rounds depends on the block |
96 | // size, key size and key words. For our purposes of random shuffle a 4 to 8 |
97 | // rounds is enough. |
98 | // W = word size |
99 | // B = 2 * W = block size |
100 | template <int B> |
101 | uint64_t index_shuffle(const uint64_t index, const std::array<uint32_t, 3>& key, |
102 | const uint64_t max_index, const int32_t rounds) { |
103 | const auto round_keys = simon_key_schedule<B / 2>(key, rounds); |
104 | uint64_t new_index = index; |
105 | while (true) { |
106 | new_index = simon_encrypt<B / 2>(new_index, round_keys); |
107 | if (new_index <= max_index) { |
108 | return new_index; |
109 | } |
110 | } |
111 | } |
112 | |
113 | #undef ROTL |
114 | #undef ROTR |
115 | #undef SIMON_F |
116 | #undef SIMON_RxC |
117 | |
118 | } // namespace impl |
119 | |
120 | uint64_t index_shuffle(const uint64_t index, const std::array<uint32_t, 3>& key, |
121 | const uint64_t max_index, const int32_t rounds) { |
122 | // Block size must be large enough to represent max_index and even (since |
123 | // word size is half of it). We force at least 16 bits as minimum block size |
124 | // since we observed pattern in the permutations below. |
125 | int block_size = static_cast<int>(std::ceil(std::log2(max_index))); |
126 | block_size = std::max(block_size + block_size % 2, kMinBlockSize); |
127 | assert(block_size > 0 && block_size % 2 == 0 && block_size <= 64); |
128 | // At least 4 rounds and number of rounds must be even. |
129 | assert(rounds >= 4 && rounds % 2 == 0); |
130 | #define HANDLE_BLOCK_SIZE(B) \ |
131 | case B: \ |
132 | return impl::index_shuffle<B>(index, key, max_index, rounds); |
133 | |
134 | switch (block_size) { |
135 | HANDLE_BLOCK_SIZE(16); |
136 | HANDLE_BLOCK_SIZE(18); |
137 | HANDLE_BLOCK_SIZE(20); |
138 | HANDLE_BLOCK_SIZE(22); |
139 | HANDLE_BLOCK_SIZE(24); |
140 | HANDLE_BLOCK_SIZE(26); |
141 | HANDLE_BLOCK_SIZE(28); |
142 | HANDLE_BLOCK_SIZE(30); |
143 | HANDLE_BLOCK_SIZE(32); |
144 | HANDLE_BLOCK_SIZE(34); |
145 | HANDLE_BLOCK_SIZE(36); |
146 | HANDLE_BLOCK_SIZE(38); |
147 | HANDLE_BLOCK_SIZE(40); |
148 | HANDLE_BLOCK_SIZE(42); |
149 | HANDLE_BLOCK_SIZE(44); |
150 | HANDLE_BLOCK_SIZE(46); |
151 | HANDLE_BLOCK_SIZE(48); |
152 | HANDLE_BLOCK_SIZE(50); |
153 | HANDLE_BLOCK_SIZE(52); |
154 | HANDLE_BLOCK_SIZE(54); |
155 | HANDLE_BLOCK_SIZE(56); |
156 | HANDLE_BLOCK_SIZE(58); |
157 | HANDLE_BLOCK_SIZE(60); |
158 | HANDLE_BLOCK_SIZE(62); |
159 | default: |
160 | return impl::index_shuffle<64>(index, key, max_index, rounds); |
161 | } |
162 | #undef HANDLE_BLOCK_SIZE |
163 | } |
164 | |
165 | } // namespace random |
166 | } // namespace tensorflow |
167 | |