From 096cc4a7c6e81e5ae60698728954099028d49a0f Mon Sep 17 00:00:00 2001
From: Mohd Tarmizi <mohtar@users.noreply.github.com>
Date: Sun, 19 Jun 2016 15:24:48 +0800
Subject: [PATCH] Recover on panic from handler

---
 src/lib.rs | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 087a0bd..c3af671 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -44,6 +44,7 @@ use std::collections::{HashMap, HashSet};
 use std::io::{self, Read, Write, Cursor, BufRead};
 use std::mem;
 use std::net::TcpListener;
+use std::panic::{AssertUnwindSafe, RefUnwindSafe, catch_unwind};
 use std::rc::Rc;
 
 #[cfg(unix)] use unix::{Transport, Socket};
@@ -561,7 +562,8 @@ impl Drop for Request {
     }
 }
 
-fn run_transport<F>(mut handler: F, transport: &mut Transport) where F: FnMut(Request) {
+fn run_transport<F>(handler: F, transport: &mut Transport) where
+        F: Fn(Request) + RefUnwindSafe {
     let addrs: Option<HashSet<String>> = match std::env::var("FCGI_WEB_SERVER_ADDRS") {
         Ok(value) => Some(value.split(',').map(|s| s.to_owned()).collect()),
         Err(std::env::VarError::NotPresent) => None,
@@ -580,10 +582,12 @@ fn run_transport<F>(mut handler: F, transport: &mut Transport) where F: FnMut(Re
             None => true,
         };
         if allow {
-            let sock = Rc::new(sock);
+            let sock = AssertUnwindSafe(Rc::new(sock));
             loop {
                 let (request_id, role, keep_conn) = Request::begin(&sock).unwrap();
-                handler(Request::new(sock.clone(), request_id, role).unwrap());
+                catch_unwind(|| {
+                    handler(Request::new(sock.clone(), request_id, role).unwrap());
+                }).unwrap_or(());
                 if !keep_conn { break; }
             }
         }
@@ -594,7 +598,7 @@ fn run_transport<F>(mut handler: F, transport: &mut Transport) where F: FnMut(Re
 /// Runs as a FastCGI process with the given handler.
 ///
 /// Available under Unix only. If you are using Windows, use `run_tcp` instead.
-pub fn run<F>(handler: F) where F: FnMut(Request) {
+pub fn run<F>(handler: F) where F: Fn(Request) + RefUnwindSafe {
     run_transport(handler, &mut Transport::new())
 }
 
@@ -603,19 +607,22 @@ pub fn run<F>(handler: F) where F: FnMut(Request) {
 /// Unix domain sockets are supported.
 ///
 /// Available under Unix only.
-pub fn run_raw<F>(handler: F, raw_fd: std::os::unix::io::RawFd) where F: FnMut(Request) {
+pub fn run_raw<F>(handler: F, raw_fd: std::os::unix::io::RawFd) where
+        F: Fn(Request) + RefUnwindSafe {
     run_transport(handler, &mut Transport::from_raw_fd(raw_fd))
 }
 
 #[cfg(unix)]
 /// Accepts requests from a user-supplied TCP listener.
-pub fn run_tcp<F>(handler: F, listener: &TcpListener) where F: FnMut(Request) {
+pub fn run_tcp<F>(handler: F, listener: &TcpListener) where
+        F: Fn(Request) + RefUnwindSafe {
     use std::os::unix::io::AsRawFd;
     run_transport(handler, &mut Transport::from_raw_fd(listener.as_raw_fd()))
 }
 
 #[cfg(windows)]
 /// Accepts requests from a user-supplied TCP listener.
-pub fn run_tcp<F>(handler: F, listener: &TcpListener) where F: FnMut(Request) {
+pub fn run_tcp<F>(handler: F, listener: &TcpListener) where
+        F: Fn(Request) + RefUnwindSafe {
     run_transport(handler, &mut Transport::from_tcp(&listener))
 }