关于使用 rust wasm 插件在处理 http body 时遇到的问题 #1465

jaymie9019 · 2024-10-31T02:24:47Z

jaymie9019
Oct 31, 2024

https://github.com/alibaba/higress/blob/main/plugins/wasm-rust/example/sse-timing/src/lib.rs

以下这段代码

    fn on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) -> DataAction {
        if !self.is_event_stream {
            return DataAction::Continue;
        }

        let body = self
            .get_http_response_body(0, body_size)
            .unwrap_or_default();
        self.event_stream.update(body);
        self.process_event_stream(end_of_stream)
    }

当 http 响应的 body 比较大的时候 end_of_stream 会显示 false，false，最后变成 true，但是变成 true 的时候 body_size 是 0 。这时候下面这段代码

    let body = self
        .get_http_response_body(0, body_size)
        .unwrap_or_default();

真的能读取到完整的数据吗？我测试下来是直接读取失败的.

为了解决上面的问题，我在写代码的时候自己写个一个 buff 来收集每次 end_of_stream = false 的时候的 body 数据

然后把这个 body 压缩之后再写入到 body 中
self.set_http_response_body(0, self.res_body_len, buffer.as_slice());

但是这个方法表现出的现象非常奇怪，压缩部分的内容是加在原来的 body 之后的。并没有覆盖掉原来的

下面是一部分代码片段

impl HttpContext for MsePlugin {

    fn on_http_response_headers(
        &mut self,
        _num_headers: usize,
        _end_of_stream: bool,
    ) -> HeaderAction {
        if self.need_zip {
            self.set_http_response_header("content-length", None);
            self.set_http_response_header("content-encoding", Some("gzip"));
            self.log.info(&format!("on_http_response_headers set content-length None, content-encoding gzip"));
        }
        HeaderAction::Continue
    }

    fn on_http_response_body(&mut self, _body_size: usize, _end_of_stream: bool) -> DataAction {
        if !self.need_zip {
            return DataAction::Continue;
        }

        let mut trace_id = String::new();
        if let Some(trace_header) = self.res_headers.get("trace_id") {
            trace_id = trace_header.clone();
        }

        self.log.info(&format!("httpContext on_http_response_body _body_size: {}, _end_of_stream: {}, trace_id: {}",
                               _body_size, _end_of_stream,
                               trace_id));
        self.res_body_len += _body_size;

        if _body_size > 0 {
            if let Some(body_bytes) = self.get_http_response_body(0, _body_size) {
                // self.log.info(&format!("on_http_response_body read success body size: {}, traceId: {}", body_bytes.len(), trace_id));
                self.resp_body_buffer.extend_from_slice(body_bytes.as_slice());
            }
        }

        if !_end_of_stream {
            return DataAction::Continue;
        }

        if let Some(config) = &self.config {
            self.log.info(&format!("on_http_response_body trace_id: {}, config: {:?}", trace_id, config));

            let mut buffer = Vec::new();
            let mut encoder = GzEncoder::new(
                self.resp_body_buffer.as_slice(),
                flate2::Compression::new(config.get_comp_level()));
            encoder.read_to_end(&mut buffer).unwrap();
            self.log.info(&format!("on_http_response_body gziped response body size: {} => {}, traceId: {}", self.resp_body_buffer.len(), buffer.len(), trace_id));
            // self.set_http_response_body(0, self.res_body_len, buffer.as_slice());
            self.replace_http_response_body(buffer.as_slice());
        }
        DataAction::Continue
    }
}

测试的效果截图，可以看到加在原来的 body 后面了

Answered by jaymie9019

Oct 31, 2024

我成功了！最终的代码

多了一行，每次把原来的 body 给擦掉

self.replace_http_response_body(Vec::new().as_slice());

    fn on_http_response_body(&mut self, _body_size: usize, _end_of_stream: bool) -> DataAction {
        if !self.need_zip {
            return DataAction::Continue;
        }

        let mut trace_id = String::new();
        if let Some(trace_header) = self.res_headers.get("trace_id") {
            trace_id = trace_header.clone();
        }

        self.log.info(&format!("httpContext on_http_response_body _body_size: {}, _end_of_stream: {}, trace_id: {}",
                               _body_size, _end_of_stream,
                               trace_id));

        // 自己缓存
        self.res_body_len +…

View full answer

jaymie9019 · 2024-10-31T06:48:35Z

jaymie9019
Oct 31, 2024
Author

@007gzs 特别想要请教下 007 大佬

0 replies

007gzs · 2024-10-31T07:08:41Z

007gzs
Oct 31, 2024
Collaborator

每次回调on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) 时，说明有 body_size 新数据可以读取，如果end_of_stream不为ture，说明后面还有数据，如果本次你return 了continue,下次来数据的时候就获取不到本次的数据了。（同时本次的数据也没法被replace了）

想一次性获取所有body数据可以参考

higress/plugins/wasm-rust/src/plugin_wrapper.rs

Lines 447 to 461 in 9a1edcd

    
           self.res_body_len += body_size; 
        
           if !end_of_stream { 
        
               return DataAction::StopIterationAndBuffer; 
        
           } 
        
           let mut res_body = Bytes::new(); 
        
           if self.res_body_len > 0 { 
        
               if let Some(body) = self.get_http_response_body(0, self.res_body_len) { 
        
                   res_body = body; 
        
               } 
        
           } 
        
           self.http_content 
        
               .borrow_mut() 
        
               .on_http_response_complete_body(&res_body)

16 replies

jaymie9019 Oct 31, 2024
Author

是进不了on_response_headers还是set_response_header无效? 你在if外面打log能打出来么

我这个问题 header 的问题解决了，还剩下最后一个问题就是大 body 时候的问题了

我现在代码是这样子的，直接 continue ，然后每次读取数据然后压缩 replace 掉。现在出现的问题是响应会被截断掉，而且截断的地方不固定。

    fn on_http_response_body(&mut self, _body_size: usize, _end_of_stream: bool) -> DataAction {
        if !self.need_zip {
            return DataAction::Continue;
        }
    
        let mut trace_id = String::new();
        if let Some(trace_header) = self.res_headers.get("trace_id") {
            trace_id = trace_header.clone();
        }
    
        self.log.info(&format!("httpContext on_http_response_body _body_size: {}, _end_of_stream: {}, trace_id: {}",
                               _body_size, _end_of_stream,
                               trace_id));
    
        if _body_size == 0 {
            return DataAction::Continue;
        }
    
        if let Some(config) = &self.config {
            if let Some(body_bytes) = self.get_http_response_body(0, _body_size) {
                // self.log.info(&format!("on_http_response_body read success body size: {}, traceId: {}", body_bytes.len(), trace_id));
                // self.resp_body_buffer.extend_from_slice(body_bytes.as_slice());
                let mut encoder = GzEncoder::new(Vec::new(), flate2::Compression::new(config.get_comp_level())
                );
    
                // 将数据写入压缩器
                let _ = encoder.write_all(body_bytes.as_slice());
                // 完成压缩，获取压缩后的数据
                let compressed_data = encoder.finish().expect("Failed to finish encoding");
                self.replace_http_response_body(compressed_data.as_slice());
            }
        }
    
        DataAction::Continue
    }

007gzs Oct 31, 2024
Collaborator

你这是每个分片单独压缩，最后组好的数据没法正常解压把。截断的原因可能就是client解压不了。

jaymie9019 Oct 31, 2024
Author

你这是每个分片单独压缩，最后组好的数据没法正常解压把。截断的原因可能就是client解压不了。

我最终的代码是这样子的，虽然没成功，但是我发现了一些现象可以 share 下

    fn on_http_response_body(&mut self, _body_size: usize, _end_of_stream: bool) -> DataAction {
        if !self.need_zip {
            return DataAction::Continue;
        }

        let mut trace_id = String::new();
        if let Some(trace_header) = self.res_headers.get("trace_id") {
            trace_id = trace_header.clone();
        }

        self.log.info(&format!("httpContext on_http_response_body _body_size: {}, _end_of_stream: {}, trace_id: {}",
                               _body_size, _end_of_stream,
                               trace_id));

        // 自己缓存
        self.res_body_len += _body_size;
        if let Some(body_bytes) = self.get_http_response_body(0, self.res_body_len) {
            self.log.info(&format!("on_http_response_body read success body size: {}, traceId: {}", body_bytes.len(), trace_id));
            self.resp_body_buffer.extend_from_slice(body_bytes.as_slice());
        }


        if !_end_of_stream {
            return DataAction::Continue;
        }

        if let Some(config) = &self.config {
            let mut encoder = GzEncoder::new(Vec::new(), flate2::Compression::new(config.get_comp_level()));

            // 将数据写入压缩器
            let _ = encoder.write_all(self.resp_body_buffer.as_slice());
            // 完成压缩，获取压缩后的数据
            let compressed_data = encoder.finish().unwrap();
            self.log.info(&format!("on_http_response_body write to body success body size: {} => {}, traceId: {}", self.resp_body_buffer.len(),
                                   compressed_data.len(), trace_id));
            self.replace_http_response_body(compressed_data.as_slice());
        }

        DataAction::Continue
    }

如果 _end_of_stream = false

针对这行代码

if !_end_of_stream {
return DataAction::Continue;
}

如果返回
StopIterationAndBuffer = 1,
StopIterationNoBuffer = 3,

大的 respnse_body 无论如何都无法打印出 _end_of_stream = true。我可以认为我自己的缓存没有办法获取到全量的数据。

如果返回 continue，
我可以观察到 _end_of_stream = true，并且读到了所有的数据，并且我也压缩了数据写入了，但是这个的问题跟最开始一样了，我并没有 replace 掉原始的 body，而是在某一段里面补充进去的

jaymie9019 Oct 31, 2024
Author

我成功了！最终的代码

多了一行，每次把原来的 body 给擦掉

self.replace_http_response_body(Vec::new().as_slice());

    fn on_http_response_body(&mut self, _body_size: usize, _end_of_stream: bool) -> DataAction {
        if !self.need_zip {
            return DataAction::Continue;
        }

        let mut trace_id = String::new();
        if let Some(trace_header) = self.res_headers.get("trace_id") {
            trace_id = trace_header.clone();
        }

        self.log.info(&format!("httpContext on_http_response_body _body_size: {}, _end_of_stream: {}, trace_id: {}",
                               _body_size, _end_of_stream,
                               trace_id));

        // 自己缓存
        self.res_body_len += _body_size;
        if let Some(body_bytes) = self.get_http_response_body(0, self.res_body_len) {
            self.log.info(&format!("on_http_response_body read success body size: {}, traceId: {}", body_bytes.len(), trace_id));
            self.resp_body_buffer.extend_from_slice(body_bytes.as_slice());
            self.replace_http_response_body(Vec::new().as_slice());
        }


        if !_end_of_stream {
            return DataAction::Continue;
        }

        if let Some(config) = &self.config {
            let mut encoder = GzEncoder::new(Vec::new(), flate2::Compression::new(config.get_comp_level()));
            let _ = encoder.write_all(self.resp_body_buffer.as_slice());
            let compressed_data = encoder.finish().unwrap();
            self.log.info(&format!("on_http_response_body write to body success body size: {} => {}, traceId: {}", self.resp_body_buffer.len(),
                                   compressed_data.len(), trace_id));
            self.replace_http_response_body(compressed_data.as_slice());
        }

        DataAction::Continue
    }

Answer selected by jaymie9019

007gzs Oct 31, 2024
Collaborator

这样没问题，只是如果是超大body会全部占用插件内存

jaymie9019 Oct 31, 2024
Author

这样没问题，只是如果是超大body会全部占用插件内存

higress 会对插件的内存做限制吗？这部分内存应该在请求结束之后会立即释放吧

007gzs Oct 31, 2024
Collaborator

是会释放，但是如果并发请求你的内存会炸

007gzs Oct 31, 2024
Collaborator

参考下这个，我没调，大概是这个意思


fn on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) -> DataAction {
        if body_size > 0{
            if let Some(body) = self.get_http_response_body(0, body_size){
                self.gz.write_all(&body).unwrap();
                let gzbuf = self.gz.get_mut();
                self.replace_http_response_body(&gzbuf);
                gzbuf.clear();
            }
        }
        if end_of_stream{
            let gzbuf = self.gz.finish().unwrap();
            self.replace_http_response_body(&gzbuf);
        }
        DataAction::Continue
}

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

关于使用 rust wasm 插件在处理 http body 时遇到的问题 #1465

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Replies: 2 comments 16 replies

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Select a reply

关于使用 rust wasm 插件在处理 http body 时遇到的问题 #1465

jaymie9019 Oct 31, 2024

Replies: 2 comments · 16 replies

jaymie9019 Oct 31, 2024 Author

007gzs Oct 31, 2024 Collaborator

jaymie9019 Oct 31, 2024 Author

007gzs Oct 31, 2024 Collaborator

jaymie9019 Oct 31, 2024 Author

jaymie9019 Oct 31, 2024 Author

007gzs Oct 31, 2024 Collaborator

jaymie9019 Oct 31, 2024 Author

007gzs Oct 31, 2024 Collaborator

007gzs Oct 31, 2024 Collaborator

jaymie9019
Oct 31, 2024

Replies: 2 comments 16 replies

jaymie9019
Oct 31, 2024
Author

007gzs
Oct 31, 2024
Collaborator

jaymie9019 Oct 31, 2024
Author

007gzs Oct 31, 2024
Collaborator

jaymie9019 Oct 31, 2024
Author

jaymie9019 Oct 31, 2024
Author

007gzs Oct 31, 2024
Collaborator

jaymie9019 Oct 31, 2024
Author

007gzs Oct 31, 2024
Collaborator

007gzs Oct 31, 2024
Collaborator