Skip to content

Commit bc988e6

Browse files
committed
Proper json string and security fix
Signed-off-by: M Q <[email protected]>
1 parent 0677a89 commit bc988e6

File tree

3 files changed

+26
-6
lines changed

3 files changed

+26
-6
lines changed

platforms/aidoc/restful_app/ai_spleen_seg_app/app.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -117,8 +117,11 @@ def run(self, *args, **kwargs):
117117
"output_files": output_files,
118118
"error_message": None,
119119
"error_code": None,
120-
"result": ai_results.model_dump_json(),
121120
}
121+
122+
# Need to use pydantic function to dump to string and then reload to dict
123+
# because for some reason direct dumping to dict did not work well
124+
callback_msg_dict["result"] = json.loads(ai_results.model_dump_json())
122125
self._status_callback(json.dumps(callback_msg_dict))
123126

124127
except Exception as e:

platforms/aidoc/restful_app/app.py

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import requests
2222
from flask import Flask, jsonify, request
23+
from flask_wtf.csrf import CSRFProtect
2324

2425
# The MONAI Deploy application to be wrapped.
2526
# This can be changed to any other application in the repository.
@@ -29,7 +30,13 @@
2930
APP_CLASS_NAME = "AISpleenSegApp"
3031

3132
# Flask application setup
32-
restful_app = Flask(__name__)
33+
app = Flask(__name__)
34+
# It is recommended to use a securely generated random string for the secret key,
35+
# and store it in an environment variable or a secure configuration file.
36+
app.config["SECRET_KEY"] = os.environ.get("FLASK_SECRET_KEY", "a-secure-default-secret-key-for-dev")
37+
csrf = CSRFProtect(app)
38+
39+
3340
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
3441

3542
# Global state to track processing status. A lock is used for thread safety.
@@ -64,8 +71,16 @@ def app_status_callback(summary: str):
6471
logging.info(f"Sending final status callback to {callback_url}")
6572
# Here you could map the summary to the expected format of the callback.
6673
# For now, we'll just forward the summary.
67-
requests.post(callback_url, data=summary, timeout=5)
74+
response = requests.post(callback_url, data=summary, timeout=5)
75+
response.raise_for_status() # for bad status codes (4xx or 5xx)
6876
logging.info("Sent final status callback.")
77+
78+
except requests.exceptions.Timeout:
79+
logging.error("The request timed out.")
80+
except requests.exceptions.ConnectionError:
81+
logging.error("A connection error occurred.")
82+
except requests.exceptions.RequestException as e:
83+
logging.error(f"An unexpected error occurred: {e}")
6984
except Exception as e:
7085
logging.error(f"Failed to send callback to {callback_url}: {e}")
7186

@@ -106,13 +121,14 @@ def app_status_callback(summary: str):
106121
logging.info("Processor is now IDLE.")
107122

108123

109-
@restful_app.route("/status", methods=["GET"])
124+
@app.route("/status", methods=["GET"])
110125
def status():
111126
"""Endpoint to check the current processing status."""
112127
return jsonify({"status": get_processing_status()})
113128

114129

115-
@restful_app.route("/process", methods=["POST"])
130+
@app.route("/process", methods=["POST"])
131+
@csrf.exempt
116132
def process():
117133
"""Endpoint to start a new processing job."""
118134
if get_processing_status() == "BUSY":
@@ -151,4 +167,4 @@ def process():
151167
args = parser.parse_args()
152168
host = args.host or os.environ.get("FLASK_HOST", "0.0.0.0")
153169
port = args.port or int(os.environ.get("FLASK_PORT", 5000))
154-
restful_app.run(host=host, port=port)
170+
app.run(host=host, port=port)

platforms/aidoc/restful_app/requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,4 @@ requests>=2.32
1414
types-requests>=2.32.0
1515
Werkzeug==2.2.3
1616
pydantic>=2.9.0
17+
Flask-WTF>=1.0

0 commit comments

Comments
 (0)