将tiny-keccak的核心部分几乎彻底重构为keccak-state的过程

八千字的流水账，基本把整个keccak-state的历程和涉及到的所有东西梳理了一次（）有什么没提到的之后再补充（）

连肝了十个小时，终于几乎彻底重构了tiny-keccak的核心部分……这玩意总算被我盘光溜了……

背景介绍：tiny-keccak是Rust社区的两大Keccak/SHA3(/(c)SHAKE/K12)实现之一。另外一大是Rust Crypto组织的，他们维护了全套的密码学算法的pure Rust实现，并且所有实现都基于一整套完善且优化的抽象。但是整一套抽象还是比较重的，对于想要轻依赖的使用者来说并不友好，而tiny-keccak是0依赖的。而且Rust Crypto那套抽象基于typenum和generic array，属于const泛型还不完善的情况下强迫症比较受不了的那种权宜之计（但没办法这是现阶段大量写trait必需的）。

我一开始就是基于tiny-keccak实现BCSP协议，逐渐发现需要很多tiny-keccak没提供的工具函数，比如简化的创建状态、输出到定长数组、流密码要用到的XOR输出、甚至跳过输出。但是tiny-keccak的内部状态没公开，导致实现这些函数得套两层娃，虽然inline了没性能损失，但强迫症怎么可能受得了。而且作为流密码使用最好在释放内存之前把内部状态清零，以防缓冲区泄漏这样的事情，而保证清零操作不被编译器优化掉需要用到atomic fence和volatile write，这个Rust Crypto已经帮我们提供了一个抽象也就是zeroize库（不得不说Rust Crypto真的搞了不少基础设施，hex-literal也是他们的），但是我们知道不能给别的crate的结构实现trait，这就不得不单开一个crate才能实现了。考虑到核心的内部状态也就三四百行的规模，以及这玩意已经三年没发新版了，我开了一个tiny-keccak-core，把它的内部状态单独拿出来公开，然后再在BCSP实现里封装cSHAKE。

一开始做的改动也就是加上了释放时zeroize，对API小的调整，但是我这种强迫症怎么可能忍得住不左改改右改改。就在这时我把目光放在了一处unsafe代码，用裸指针写的XOR，对于现代编译器优化真的有必要吗？于是我上Rust Users问了一下（https://users.rust-lang.org/t/93119/），得到的结论是没必要，改成直接的for或者zip就可以，用for的话可以通过reslicing保证长度一致从而提高性能（这个问题还有后话）。

接下来就一发不可收拾了。先是看那个宏实现的Keccak函数不爽，最后直接把宏拆了，只有类型化实现用宏，把函数本身解放出来。接下来又看向encode len实现，觉得可以用ctlz优化一下，本来是没看懂瞎猫碰死耗子式优化，后来试着把测试用例写成hex才看明白，原来就是把长度的长度也做了一个编码，于是用absorb两次而不需要中间缓存的方式实现了，（之前还担心absorb两次会不会多出开销，经过后面重构fold才完全明白新写入的数据量大于rate才会置换，不过多一次调用开销也算开销），顺便直接做成了方法。为了能实现K12的完备性把delim作为了公开字段。

后来到了我在高铁上把BCSP的握手写了个大概（基本只是理清了逻辑，还完全没打磨抽象）的时候，BCSP那边的cSHAKE封装已经很完善了，除了上面提到的工具函数之外还有一个类型化的custom string实现，方便使用的同时还能保证不同custom string的实例不放错地方（这算是Rust中除了const泛型之外另一种实现dependent type的方法）。于是没怎么犹豫就把这套cSHAKE实现搬到了此时已经改名为keccak-core的仓库，BCSP那边只留一个实现custom string的宏调用。

决定让BCSP走在无TLS并且默认无HTTP握手的WebSocket上之后，看向21年写LiveKit的时候就已经很熟悉的tungstenite库，发现依赖了一个rand库来实现生成随机mask（实际上还有HTTP握手用的随机字符串）。实际上mask是为了防止跟HTTP请求混淆，对于无HTTP握手的情况，mask是没有用的，但毕竟是WebSocket的RFC里面的东西，config里加个开关还好，怎么可能给你加个feature把rand依赖关掉。而rand本身不小不说，里面的默认算法是ChaCha12，造BCSP就是为了做少依赖实验，能被Keccak替代的算法我怎么可能忍住让它出现在协议里。于是我想到了我还没把25519研究明白（现在至少明白到能自己基于base库实现密钥交换和签名了，尽管除了我自己想出来的基于交换的签名之外，ed25519风格的签名还没完全弄明白，之后可能单独讲讲25519的事，预告一下这玩意还有个448bit的版本）的时候看ed25519库如何强行替换掉SHA2-512算法的方式（https://github.com/dalek-cryptography/ed25519-dalek/issues/64），就是狸猫换太子，用cargo的patch放一个自己版本的同名包，然后里面写上自己的实现，只要调用它的库调用的方式是固定的就没问题。于是就得学着rand库实现一个thread local的状态，它直接用Rc<UnsafeCell<State>>来绕开的LocalKey只能提供不可变引用，中间害怕自己实现的不sound还去问了问（https://users.rust-lang.org/t/94278）（确实是sound的，只要封装起来不要允许直接访问引用就没问题，LocalKey不给可变引用是怕同时拿两个可变引用）。实现的时候顺便抽象出了Absorb Squeeze Once三个trait，把工具函数全封到了trait里，然后控制thread local状态只能squeeze。

实现完了rand之后，把getrandom的输入换成了可选的MaybeUninit，然后就想squeeze到定长数组这些是不是也能改成直接写到未初始化，过程中了解了一些关于未初始化内存的知识（内存是知道自己有没有被初始化的，而如果没有，那编译器想怎么来就怎么来，并不一定是指向一个固定的随机内存）（https://www.ralfj.de/blog/2019/07/14/uninit.html），看了一些资料和以往提问后直觉觉得不太可能sound的实现出来，结果问了一下确实（https://users.rust-lang.org/t/94321），而且以0初始化的情况下会有一个底层指令去找zeroed的未分配内存，所以实际上一般也没这个写开销。这个时候突然注意到一点：

（以上是25号上午10点多肝完之后写的，我现在一般三四点最晚六点睡所以对我来说算熬夜了，写到最关键的地方就倒下了，幸亏没把作息打乱了，然后以下是这两天断断续续写完的（主要是经常打断），尽可能靠笔记和录屏回忆起来比较完整的思考过程，亲身体验查资料写代码的时候做笔记和录屏都是好习惯，实际上很多描述都是基于最后完全弄明白之后的理解写的，前面一点的步骤当时可能还没理解那么清楚）

实际上squeeze的过程就是分很多次把数据从内部状态里用copy_from_slice复制出来（后面会把本质讲得更清楚），那么只要把这个复制换成XOR，不就可以在完全不需要分配的情况下实现XOR输出了吗？换成什么都不做不就能实现跳过输出了吗？理论上甚至可以这样实现从一个内部状态squeeze到另一个内部状态，这是很有用的，比如一个状态生成密钥，生成的密钥喂给加解密和MAC状态，BCSP里就有这样的结构，不过实现会比较复杂而且容易bug，所以到现在也还没实现。于是我基于他原本的Buffer接口，参考xorin和setout写了xorout和skipout，并且把state中absorb和squeeze的过程通过宏抽象出来，在调用宏时指定调用的Buffer操作方法，然后基于宏实现absorb squeeze squeeze_xor squeeze_skip方法。但原本的Buffer接口这么改有一个小问题，skipout实际上不做任何操作，所以传进去的也是一个空的slice，但是在squeeze过程中还是会reslice这个slice，而reslice空的slice不用说肯定panic。于是我做了一个小改动，把输入/输出的offset也传进xorin/{set, xor, skip}out这些方法里，然后其余的在里面reslicing，skipout什么都不做。而对于cSHAKE抽象，则把Squeeze trait拆分成Squeeze SqueezeXor SqueezeSkip，删除基于squeeze申请新内存的实现，改为直接调用state的squeeze_xor和squeeze_skip方法。

其实他这个Buffer接口是我一直看着不顺眼。Keccak的内部缓冲区需要通过两种方式读写，一种是以64位为一个单位（word）的Keccak置换（实际上是对于1600是64位一个word，对于800是32位一个word，对于最小的25甚至是1位一个word），另一种是以字节为单位的输入和输出。而这里面就涉及到一个大端序和小端序的问题，因为大端序和小端序的数字在内存中排布不一样，所以就导致大端序机和小端序机做相同的置换，然后以字节形式输出的内容是不同的。解决这个问题的办法是每次涉及切换不同类型的操作之前和之后swap一下，而内部状态本身顺理成章的是跟Keccak置换一致的u64数组，所以他们在发现这个问题之后，顺理成章的选择了封装出一个Buffer，然后把所有对字节的操作封装成execute，接收要操作字节段的offset和len以及操作的闭包（原本也就是XOR输入和复制输出，我加上了XOR输出和什么都不做），然后在小端机器上直接调用函数操作，小端机器上在字节操作之前和之后把字节操作涉及到的那些u64 swap一次，这样实现了把大端小端机的差异做了封装，很顺理成章，但是总感觉有点难受，尤其是pad操作只动两个字节还要调用两次execute。

而Rust Crypto那边的SHA3/(c)SHAKE实现我之前一直没怎么细看，只是在重实现encode len的时候看了一下（实际上encode len的实现都没区别，C的XKCP那边也是那么实现的，我的方法虽然直接但是算是新方法），只知道似乎是通过只接收长度已经是8的倍数的输入/输出实现的，相当于把byte操作也转换成了word操作，所以不需要swap endian，但是一直没看懂具体怎么实现的。为了解决Buffer的问题（也不能说是问题，就是看着不顺眼），我就想一定要研究明白，还展开了宏方便跳转。结果这次注意到Rust Crypto的实现似乎好简单啊，内部状态就u64数组的buffer、RC和delim，而RC实际上完全没用到。（这同时也引起了我对哪些内部状态可以提取成const泛型的思考。）而tiny-keccak的版本有buffer、offset（对啊为什么Rust Crypto没有buffer的offset呢）、rate、delim、mode（这个Rust Crypto是单独的squeeze reader所以不需要区分）、还有一个不占大小的phantom type用来放（类似于我cSHAKE实现中custom string的）Keccak-f和Keccak-p的类型化（实际上这个也能替换成const泛型……这么一说的话custom string也可以诶，（正在写这篇文章的时候）试了一下效果还不错，但就是目前还不支持&'static作为const泛型所以暂时是不可以）。

于是我开始仔细看Rust Crypto的实现。首先发现它的内部状态absorb进去的都是固定长度的block，而这个block的长度是……136？似乎没见过这个数字，然而很快就反应过来这是bits_to_rate(256)。为了generic array直接把rate硬编码了，也没有任何注释。也就是每次处理一block的输入，这一block也就是一个rate。联想到rate是跟安全参数挂钩的参数，安全参数的bits越大rate就越小，rate越小相当于进行置换的输入间隔就越小，安全系数越高就越需要频繁的置换，嗯，合理。而Rust Crypto的抽象直接就是update blocks，那么它是怎么处理任意长度的输入的呢？很快绕进了抽象的迷宫里，算了还是继续回去改吧。

抽离出absorb和squeeze的过程之后，我发现它们的逻辑几乎完全一致，并且absorb标注是first foldp，squeeze标注是second foldp，foldp大概是标准文档里的命名吧。于是我很快将两个宏合并为一个foldp。下一步开始处理Buffer，首先把Buffer的方法全部转移到直接在状态上，以方便之后慢慢合并。在去除Buffer的过程中我开始好奇是什么时候把buffer这玩意引入的，就在tiny-keccak原仓库blame了一下，结果发现就是在发现大小端问题之后加的（前面细说过了）。还有一个额外发现：之前那个裸指针写的XOR是来自一个PR，最早的代码就是很直接的zip，然而那个PR据当时的作者们说性能提升了35%-50%！不过那个PR是7年前的事了。立即去前面提到的帖子里继续问，是不是编译器在这7年间在这方面有更多的优化，结果得到了肯定的答复，zip的性能大幅度优化了，现在用带reslicing的for和zip都可以。

回去接着看Rust Crypto的实现。实际上Rust Crypto实现的Keccak置换函数是一个单独的crate，也没有那些复杂的抽象，而且他们通过抽象出LaneSize也就是前面提到的每个word的长度，实现了从u8对应的200到u64对应的1600的不同容量的Keccak函数，而且还提供了ARM的Keccak指令集上的实现（没错目前唯一实现了硬件加速Keccak/SHA3的主流处理器架构是ARM，Intel你给劲啊，AMD你得好好劝劝啊），所以既然已经有完善的能满足需求的实现，置换函数之后就打算用他们的实现了，这也是最后选择改名为keccak-state的原因，相当于在置换函数实现的基础上实现了一个内部状态。另外那个时候才注意到Keccak-p的RC其实就是Keccak-f的后12轮的RC，之前一直没有注意到。

接下来我准备从头开始，对照着两家的实现来重新实现。一开始打算把LaneSize也作为泛型参数，从而直接使得这个内部状态能作为从200到1600各种容量的Keccak函数的内部状态，但后面发现zeroize等很多地方麻烦比较多，就放弃了，还是只做实际使用的1600。rate和delim当时想着都可以作为const泛型。

简单搭好架子之后，接着看Rust Crypto的那套抽象是怎么处理任意长度输入的。结果发现包装器里其实还有个Buffer，这个Buffer里除了有一个block大小的的数组以外还有一个offset，这下就知道offset去哪了……而且他那个offset还是u8的，虽然block长度小于256，但其实没什么意义，反正每次用的时候都得转成usize，看似省了空间，实际上每次都得还原回同样多的空间，还多了个转换的开销呢。然后定位到一个叫digest block的方法，接收一个任意长度的输入和一个接收若干block去改变（具体算法）内部状态的闭包。实际上这个东西就基本上对应tiny-keccak那边的foldp了，只不过那边没有单独的block buffer。里面的逻辑大概就是把输入的长度和block长度和offset分了几种情况讨论，比如输入的长度还不到当前buffer剩余的长度，那就只复制进来不调用改变内部状态，再比如输入比block长，那就用chunk分为和block等长的好几段，再逐个调用改变内部状态，然后有还剩下的就再放到buffer里面，等下次更多数据来了一块进去。这个Buffer的实现似乎还采取了一些精巧的手段来给编译器更多的信息，来防止生成panic path，总之设计都很精致。而swap endian的问题，除了输入输出都是一block长之外（输出也是差不多的过程），实际上是输入输出都是直接用8个字节chunk，然后输入的时候每8个字节from_le_bytes，输出的时候每8个字节to_le_bytes，在这个过程中自然地完成转换，跟内部状态直接接触的都是正确endian的u64。

在完全弄明白这些之后，觉得其实那堆抽象也就是那么回事，而且不用怎么想就能知道这中间引入了多少开销：在内部状态之外多了一个block buffer的开销，在内部状态和block buffer之间相互复制的开销，转换u64可能的复制开销……还是简单直接的tiny-keccak更适合我。不过它的处理方式引起了我的思考：内部状态必须是u64吗？必须是字节操作的时候transmute再swap endian吗？能不能反过来，内部状态本身是字节，字节操作的时候就直接来，而置换操作的时候transmute再swap endian？因为字节操作很多时候只操作一部分甚至只操作一个字节，所以就得选择一部分被操作了的swap endian，而置换操作的时候一定是操作整个状态，所以整个swap就可以。为什么已经决定把字节操作和置换操作放一起了，还要让碎片化、频繁的字节操作迁就整块的置换操作？搞得pad两个字节也得调用一下execute，对于大端机swap四次？虽然几乎没有大端机，但是这河里吗？这一点都不河里。

于是放弃之前写了个开头的从零开始，继续拆execute。为了方便称呼内部状态的buf和外部输入输出的buf，将外部输入输出的buf统一称为iobuf，然后统一一下absorb squeeze squeeze_xor squeeze_skip几个方法对iobuf reslicing的过程。接下来就是把内部状态的buf改成字节数组，然后在把transmute和swap endian这些都放进置换方法里，execute只保留最简单的reslicing和调用。

接下来先做的改动是把Keccak-f和Keccak-p的类型化删掉，改成const泛型。由于const泛型不支持自定义类型，所以只能放一个bool，用true代表Keccak-f，false代表Keccak-p，并且写好常量，真是C-style啊，不过比多一个phantom data舒服。下一步继续加const泛型，把rate也作为const泛型，能给每个内部状态省8byte的空间呢。相同的办法，定义一堆R128 R224 R256 R384 R512的常量来装这些rate值。接下来想把delim也作为const泛型，但是问题有点多，因为在很多算法里delim是变化的，比如cSHAKE在name和custom string都为空的时候就是SHAKE，用的是SHAKE的delim，K12更是有好几种delim，再加上delim就一个字节，所以还是算了。不过还是把以前直接把delim公开的方式优化了一下，改成了change delim方法返回重新构建的改了delim的内部状态。

下一步就是拆除只剩下reslicing和调用的execute了，把reslicing拆到xorin/{set, xor, skip}out几个方法里，这样它们就是整齐的对buf和iobuf reslicing然后调用接收一个&mut [u8]一个&[u8]的方法了，pad也终于恢复了只是XOR两个字节的方式。

接下来就是重量级中的重量级了，把foldp抽象为一个方法。最后选择的接口形式是，传入一个要输入/输出字节（iobuf）的长度，和一个用于操作iobuf的闭包，这个闭包会传入整个内部状态的buf、内部状态的offset、iobuf的offset和本次操作的len，然后在这个闭包里面reslicing然后操作。然后基于foldp传入xor(buf, iobuf)的操作函数就是absorb，传入copy(iobuf, buf)就是squeeze，传入xor(iobuf, buf)就是squeeze_xor，传入什么也不做就是squeeze_skip。接口上有点像原先的execute，但之前的是闭包带着内部状态buf操作iobuf，现在的是闭包带着iobuf操作内部状态buf，刚好反过来。这样能带来两个好处，一个是解决了输入是不可变引用、输出是可变引用、skip甚至是没有引用的差异问题，把iobuf带在闭包里而不是直接传进去，使得给几个输入输出方法提供统一的foldp接口成为可能，特别是对于skip，不仅没有不能reslicing一个空的slice的问题，而且连内部状态buf都不需要reslicing了（虽然应该会优化掉）；另一个是对于for实现的XOR来说，reslicing就在旁边，而不是分散在两个函数离得八丈远，不需要单独抽象出来的函数来做重新告诉编译器这俩一样长的工作。极度舒适，没有一行冗余代码，既不需要一个u64单位的内部状态跟一个字节单位的block buffer然后之间复制来复制去，也不需要只有一个u64的内部状态然后每次字节操作哪怕只是动一个字节都得swap一遍endian，真正符合了duplex清洁高效的特点。强迫症得到完全满足，即使当时已经肝了十个小时也倍感欣慰。另外在改出这个新的foldp抽象的过程中，我也把原本简写的变量名全部改成最明确的，改完了就发现自己大概看懂了。实际上跟Rust Crypto那边的absorb block做的事情基本上是一样的，相当于也是把输入的长度和block长度和offset分了几种情况讨论，只不过纯用几个变量做的判断，带了几个数字试了试，逻辑是没有问题的。

另外在完全理解这个过程（应该对应sponge-duplex的duplex部分，毕竟符合没有输出转换的特点）之后，其实有个很好的比喻方式。整个过程就像是种地，xorin就是把输入数据播种到内部状态，Keccak置换就是耕地，{set, xor, skip}不管哪种out就是从内部状态收获输出数据。播种的长度是任意的，收获的长度也是任意的，反正种/收够了rate那么多就翻一次土。要求的安全级别越高，翻土就越频繁。

下一步执行之前做的决定，换成Rust Crypto的Keccak置换函数实现，顺便把字节形式和u64形式的元素数量都通过容量bits+const函数定义了，这样方便未来直接通过LaneSize或者const泛型选择容量。

这个时候已经没什么要改的了，项目改名为keccak-state，调整一下项目结构的遗留问题，当天就去写这篇的前半部分内容了。

然后是今天接着写的时候提到custom string能不能用const泛型，试了一下不行，然后写开了就接着写，先实现了cSHAKE状态的reset。cSHAKE是有rate、name、custom string这几个初始状态的，所以reset直接把内部状态buf清空是不行的。Rust Crypto那边的实现方式是多一个reset feature，开启的时候在内部状态里面加一个填了几个初始状态之后的状态，reset的时候复制一份。我的cSHAKE实现不需要这样做，因为custom string是类型化实现，所以随时能访问到，reset的时候再填一遍初始状态就好了。然后由reset想到reseed，前面说到我学着rand库做了一个thread local的用于生成随机数据的状态，而rand库的实现是会自动reseed的，每输出了指定数量的字节就从系统重新获取一次熵。于是我也写了一个简单的reseedable状态，在cSHAKE状态的基础上加一个remain，每次squeeze完了检查一下有没有用完。另外在通用的cSHAKE的状态上也加了一个通过feature开启的seed方法。很快就想到一个问题：万一一次要的字节数量就超过了reseed间隔怎么办？先不说怎么办，减掉之后remain就成负的了。尝试思考了一会之后，我突然意识到实际上这是跟foldp类似的问题，输入/输出的长度和内部状态的长度和offset可能存在各种情况，只不过是rate换成了reseed间隔，执行置换变成了执行reseed。于是我把foldp直接套上去改了一下，结果带数字看逻辑完全没问题。感觉或许有必要把foldp抽象成单独的逻辑了，就跟Rust Crypto的block buffer抽象一样，但不需要额外的buffer。另外也调整了一下foldp的过程，他原本的foldp把buf_offset弄了一个内部变量，实际上直接用内部状态的字段self.offset就可以。每次不管是置换还是reseed都一定会伴随着offset的归零，所以我就直接把offset归零放到置换/reseed方法里了。另外还有一个涉及到置换和offset归零的操作，就是cSHAKE填完初始状态之后的fill_block，实际上现在也好理解了，就是无论你offset在哪都置换一次并且offset归零，相当于多输入了这个rate长度（也就是这个block）剩下的长度个0（XOR 0等于没变），也就是把当前这个rate长度（这个block）“填满了”（0）。

然后就又由reset的问题想到之前就想过的预计算初始状态，因为正常情况下一个系统（比如说BCSP）有很多cSHAKE custom，每个新的custom状态都需要absorb一遍相同的初始状态，对于要初始化很多状态的系统来说，这无疑是不容忽视的性能开销。于是实际上可以在custom的类型化trait中加一个常量，来放置这个custom的初始状态，每次创建或者reset的时候都直接复制这个初始状态就可以。然后这个初始状态可以用构建脚本或者过程宏计算，const上下文和macro rules是肯定不能做这么复杂的计算的，之后还得写一个过程宏。我也试了一下生成出来粘贴进代码里，结论是基本不可行，每个custom 200字节呢，放在源码里还是比较碍眼的。添加了内部状态的从初始状态生成的方法和输出内部状态的方法（生成的时候用）。顺便补充了一下custom string的类型化实现，补上了可选的name，增加了直接判断是否两者都为空和直接返回delim的方便的函数。

（其实这后半部分大部分也是（对我来说的）熬到上午写的，前面要么各种事打断要么写不出来，等到7点钟才进入状态，然后写到现在11点，希望作息能跟三天前一样不乱掉吧，主要是不写完心里就会一直有这件事，没法开别的工，所以只好憋一口气无论如何写完）