ark_core/
extension.rs

1use bitcoin::opcodes::all::OP_RETURN;
2use bitcoin::script::Instruction;
3use bitcoin::Amount;
4use bitcoin::Script;
5use bitcoin::ScriptBuf;
6use bitcoin::TxOut;
7
8pub const MAGIC_BYTES: [u8; 3] = [0x41, 0x52, 0x4b];
9
10#[derive(Debug, thiserror::Error)]
11pub enum ExtensionError {
12    #[error("extension payload length overflows")]
13    PayloadLengthOverflow,
14    #[error("truncated extension packet type")]
15    TruncatedPacketType,
16    #[error("truncated extension packet length")]
17    TruncatedPacketLength,
18    #[error("truncated extension packet payload, expected {expected} bytes got {got}")]
19    TruncatedPacketPayload { expected: usize, got: usize },
20    #[error("duplicate extension packet type {0}")]
21    DuplicatePacketType(u8),
22}
23
24pub fn encode_uvarint(buf: &mut Vec<u8>, mut value: u64) {
25    loop {
26        let mut byte = (value & 0x7f) as u8;
27        value >>= 7;
28        if value != 0 {
29            byte |= 0x80;
30        }
31        buf.push(byte);
32        if value == 0 {
33            break;
34        }
35    }
36}
37
38fn decode_uvarint(data: &[u8], offset: &mut usize) -> Result<u64, ExtensionError> {
39    let mut value = 0_u64;
40    for shift in (0..64).step_by(7) {
41        let Some(byte) = data.get(*offset) else {
42            return Err(ExtensionError::TruncatedPacketLength);
43        };
44        *offset += 1;
45        value |= u64::from(byte & 0x7f) << shift;
46        if byte & 0x80 == 0 {
47            return Ok(value);
48        }
49    }
50    Err(ExtensionError::PayloadLengthOverflow)
51}
52
53pub fn is_extension(script: &Script) -> bool {
54    extension_payload(script).is_some()
55}
56
57pub fn extension_payload(script: &Script) -> Option<&[u8]> {
58    let mut instructions = script.instructions();
59    if !matches!(instructions.next(), Some(Ok(Instruction::Op(OP_RETURN)))) {
60        return None;
61    }
62    let Some(Ok(Instruction::PushBytes(bytes))) = instructions.next() else {
63        return None;
64    };
65    let bytes = bytes.as_bytes();
66    (bytes.len() >= MAGIC_BYTES.len() && bytes[..MAGIC_BYTES.len()] == MAGIC_BYTES).then_some(bytes)
67}
68
69pub fn iter_packets(payload: &[u8]) -> Result<Vec<(u8, &[u8])>, ExtensionError> {
70    let mut packets = Vec::new();
71    let mut offset = MAGIC_BYTES.len();
72
73    while offset < payload.len() {
74        let Some(packet_type) = payload.get(offset).copied() else {
75            return Err(ExtensionError::TruncatedPacketType);
76        };
77        offset += 1;
78
79        let packet_len = decode_uvarint(payload, &mut offset)? as usize;
80        let end = offset
81            .checked_add(packet_len)
82            .ok_or(ExtensionError::PayloadLengthOverflow)?;
83        if end > payload.len() {
84            return Err(ExtensionError::TruncatedPacketPayload {
85                expected: end,
86                got: payload.len(),
87            });
88        }
89
90        packets.push((packet_type, &payload[offset..end]));
91        offset = end;
92    }
93
94    Ok(packets)
95}
96
97pub fn find_packet_payload(
98    tx: &bitcoin::Transaction,
99    packet_type: u8,
100) -> Result<Option<&[u8]>, ExtensionError> {
101    for output in &tx.output {
102        let Some(payload) = extension_payload(&output.script_pubkey) else {
103            continue;
104        };
105        for (current_type, current_payload) in iter_packets(payload)? {
106            if current_type == packet_type {
107                return Ok(Some(current_payload));
108            }
109        }
110        return Ok(None);
111    }
112
113    Ok(None)
114}
115
116pub fn packet_txout(packet_type: u8, packet_payload: &[u8]) -> TxOut {
117    let mut payload = Vec::new();
118    payload.extend_from_slice(&MAGIC_BYTES);
119    push_packet(&mut payload, packet_type, packet_payload);
120
121    TxOut {
122        value: Amount::ZERO,
123        script_pubkey: op_return_script(&payload),
124    }
125}
126
127pub fn add_packet_to_psbt(
128    psbt: &mut bitcoin::Psbt,
129    packet_type: u8,
130    packet_payload: &[u8],
131) -> Result<(), ExtensionError> {
132    let mut encoded_packet = Vec::new();
133    push_packet(&mut encoded_packet, packet_type, packet_payload);
134
135    for output in &mut psbt.unsigned_tx.output {
136        let Some(existing_payload) = extension_payload(&output.script_pubkey) else {
137            continue;
138        };
139
140        for (existing_type, _) in iter_packets(existing_payload)? {
141            if existing_type == packet_type {
142                return Err(ExtensionError::DuplicatePacketType(packet_type));
143            }
144        }
145
146        let mut payload = existing_payload.to_vec();
147        payload.extend_from_slice(&encoded_packet);
148        output.script_pubkey = op_return_script(&payload);
149        return Ok(());
150    }
151
152    let txout = packet_txout(packet_type, packet_payload);
153    let len = psbt.unsigned_tx.output.len();
154
155    if len == 0 {
156        psbt.unsigned_tx.output.push(txout);
157        psbt.outputs.push(bitcoin::psbt::Output::default());
158        return Ok(());
159    }
160
161    let anchor_index = len - 1;
162    psbt.unsigned_tx.output.insert(anchor_index, txout);
163    psbt.outputs
164        .insert(anchor_index, bitcoin::psbt::Output::default());
165    Ok(())
166}
167
168fn push_packet(payload: &mut Vec<u8>, packet_type: u8, packet_payload: &[u8]) {
169    payload.push(packet_type);
170    encode_uvarint(payload, packet_payload.len() as u64);
171    payload.extend_from_slice(packet_payload);
172}
173
174fn op_return_script(data: &[u8]) -> ScriptBuf {
175    let mut script = Vec::new();
176    script.push(OP_RETURN.to_u8());
177    push_data(&mut script, data);
178    ScriptBuf::from_bytes(script)
179}
180
181fn push_data(script: &mut Vec<u8>, data: &[u8]) {
182    let len = data.len();
183    if len <= 75 {
184        script.push(len as u8);
185    } else if len <= 0xff {
186        script.push(0x4c);
187        script.push(len as u8);
188    } else if len <= 0xffff {
189        script.push(0x4d);
190        script.extend_from_slice(&(len as u16).to_le_bytes());
191    } else {
192        script.push(0x4e);
193        script.extend_from_slice(&(len as u32).to_le_bytes());
194    }
195    script.extend_from_slice(data);
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use bitcoin::absolute;
202    use bitcoin::transaction;
203    use bitcoin::TxIn;
204
205    #[test]
206    fn encodes_go_uvarint_packet_lengths() {
207        let txout = packet_txout(0x01, &[0; 136]);
208        let payload = extension_payload(&txout.script_pubkey).unwrap();
209        assert_eq!(&payload[..5], &[0x41, 0x52, 0x4b, 0x01, 0x88]);
210        assert_eq!(payload[5], 0x01);
211    }
212
213    #[test]
214    fn appends_to_existing_extension_output() {
215        let mut psbt = bitcoin::Psbt::from_unsigned_tx(bitcoin::Transaction {
216            version: transaction::Version::TWO,
217            lock_time: absolute::LockTime::ZERO,
218            input: vec![TxIn::default()],
219            output: vec![
220                TxOut {
221                    value: Amount::ZERO,
222                    script_pubkey: packet_txout(0x00, &[0xaa]).script_pubkey,
223                },
224                TxOut {
225                    value: Amount::ZERO,
226                    script_pubkey: ScriptBuf::new(),
227                },
228            ],
229        })
230        .unwrap();
231
232        add_packet_to_psbt(&mut psbt, 0x01, &[0xbb, 0xcc]).unwrap();
233
234        assert_eq!(psbt.unsigned_tx.output.len(), 2);
235        let payload = extension_payload(&psbt.unsigned_tx.output[0].script_pubkey).unwrap();
236        let packets = iter_packets(payload).unwrap();
237        assert_eq!(
238            packets,
239            vec![(0x00, &[0xaa][..]), (0x01, &[0xbb, 0xcc][..])]
240        );
241    }
242}
ark_core/extension.rs

ark_core/
extension.rs
