Merge pull request #43639 from mhugent/network_request_preprocessor

Network request preprocessor
qgis · Jun 24, 2021 · cb2fb67 · cb2fb67
2 parents 23b4786 + 408ed54
commit cb2fb67
Show file tree

Hide file tree

Showing 5 changed files with 191 additions and 6 deletions.
diff --git a/python/core/auto_generated/network/qgsnetworkaccessmanager.sip.in b/python/core/auto_generated/network/qgsnetworkaccessmanager.sip.in
@@ -321,6 +321,65 @@ The contents of the reply will be returned after the request is completed or an
 .. versionadded:: 3.6
 %End
 
+    static QString setRequestPreprocessor( SIP_PYCALLABLE / AllowNone / );
+%Docstring
+Sets a request pre-processor function, which allows manipulation of a network request before it is processed.
+
+The ``processor`` function takes the QNetworkRequest as its argument, and can mutate the request if necessary.
+
+:return: An auto-generated string uniquely identifying the preprocessor, which can later be
+         used to remove the preprocessor (via a call to :py:func:`~QgsNetworkAccessManager.removeRequestPreprocessor`).
+
+.. seealso:: :py:func:`removeRequestPreprocessor`
+
+.. versionadded:: 3.22
+%End
+%MethodCode
+    PyObject *s = 0;
+    Py_BEGIN_ALLOW_THREADS
+    Py_XINCREF( a0 );
+    QString id = QgsNetworkAccessManager::setRequestPreprocessor( [a0]( QNetworkRequest *arg )->QString
+    {
+      QString res;
+      SIP_BLOCK_THREADS
+      PyObject *s = sipCallMethod( NULL, a0, "D", arg, sipType_QNetworkRequest, NULL );
+      int state;
+      int sipIsError = 0;
+      QString *t1 = reinterpret_cast<QString *>( sipConvertToType( s, sipType_QString, 0, SIP_NOT_NONE, &state, &sipIsError ) );
+      if ( sipIsError == 0 )
+      {
+        res = QString( *t1 );
+      }
+      sipReleaseType( t1, sipType_QString, state );
+      SIP_UNBLOCK_THREADS
+      return res;
+    } );
+
+    s = sipConvertFromNewType( new QString( id ), sipType_QString, 0 );
+    Py_END_ALLOW_THREADS
+    return s;
+%End
+
+    static void removeRequestPreprocessor( const QString &id );
+%Docstring
+Removes the custom pre-processor function with matching ``id``.
+
+The ``id`` must correspond to a pre-processor previously added via a call to :py:func:`~QgsNetworkAccessManager.setRequestPreprocessor`.
+
+Returns ``True`` if processor existed and was removed.
+
+.. seealso:: :py:func:`setRequestPreprocessor`
+
+.. versionadded:: 3.22
+%End
+%MethodCode
+    if ( !QgsNetworkAccessManager::removeRequestPreprocessor( *a0 ) )
+    {
+      PyErr_SetString( PyExc_KeyError, QStringLiteral( "No processor with id %1 exists." ).arg( *a0 ).toUtf8().constData() );
+      sipIsErr = 1;
+    }
+%End
+
     void requestAuthOpenBrowser( const QUrl &url ) const;
 %Docstring
 Forwards an external browser login ``url`` opening request to the authentication handler.
@@ -356,6 +415,14 @@ Abort any outstanding external browser login request.
 
 
 
+    void preprocessRequest( QNetworkRequest *req ) const;
+%Docstring
+Preprocesses request
+
+:param req: the request to preprocess
+
+.. versionadded:: 3.22
+%End
 
   signals:
 

diff --git a/src/core/network/qgsnetworkaccessmanager.cpp b/src/core/network/qgsnetworkaccessmanager.cpp
@@ -43,6 +43,7 @@
 #include <QThreadStorage>
 #include <QAuthenticator>
 #include <QStandardPaths>
+#include <QUuid>
 
 #ifndef QT_NO_SSL
 #include <QSslConfiguration>
@@ -53,6 +54,8 @@
 
 QgsNetworkAccessManager *QgsNetworkAccessManager::sMainNAM = nullptr;
 
+static std::vector< std::pair< QString, std::function< void( QNetworkRequest * ) > > > sCustomPreprocessors;
+
 /// @cond PRIVATE
 class QgsNetworkProxyFactory : public QNetworkProxyFactory
 {
@@ -337,6 +340,11 @@ QNetworkReply *QgsNetworkAccessManager::createRequest( QNetworkAccessManager::Op
     pReq->setAttribute( QNetworkRequest::CacheSaveControlAttribute, false );
   }
 
+  for ( const auto &preprocessor :  sCustomPreprocessors )
+  {
+    preprocessor.second( pReq );
+  }
+
   static QAtomicInt sRequestId = 0;
   const int requestId = ++sRequestId;
   QByteArray content;
@@ -806,6 +814,31 @@ QgsNetworkReplyContent QgsNetworkAccessManager::blockingPost( QNetworkRequest &r
   return br.reply();
 }
 
+QString QgsNetworkAccessManager::setRequestPreprocessor( const std::function<void ( QNetworkRequest * )> &processor )
+{
+  QString id = QUuid::createUuid().toString();
+  sCustomPreprocessors.emplace_back( std::make_pair( id, processor ) );
+  return id;
+}
+
+bool QgsNetworkAccessManager::removeRequestPreprocessor( const QString &id )
+{
+  const size_t prevCount = sCustomPreprocessors.size();
+  sCustomPreprocessors.erase( std::remove_if( sCustomPreprocessors.begin(), sCustomPreprocessors.end(), [id]( std::pair< QString, std::function< void( QNetworkRequest * ) > > &a )
+  {
+    return a.first == id;
+  } ), sCustomPreprocessors.end() );
+  return prevCount != sCustomPreprocessors.size();
+}
+
+void QgsNetworkAccessManager::preprocessRequest( QNetworkRequest *req ) const
+{
+  for ( const auto &preprocessor :  sCustomPreprocessors )
+  {
+    preprocessor.second( req );
+  }
+}
+
 
 //
 // QgsNetworkRequestParameters

diff --git a/src/core/network/qgsnetworkaccessmanager.h b/src/core/network/qgsnetworkaccessmanager.h
@@ -516,6 +516,71 @@ class CORE_EXPORT QgsNetworkAccessManager : public QNetworkAccessManager
      */
     static QgsNetworkReplyContent blockingPost( QNetworkRequest &request, const QByteArray &data, const QString &authCfg = QString(), bool forceRefresh = false, QgsFeedback *feedback = nullptr );
 
+    /**
+     * Sets a request pre-processor function, which allows manipulation of a network request before it is processed.
+     *
+     * The \a processor function takes the QNetworkRequest as its argument, and can mutate the request if necessary.
+     *
+     * \returns An auto-generated string uniquely identifying the preprocessor, which can later be
+     * used to remove the preprocessor (via a call to removeRequestPreprocessor()).
+     *
+     * \see removeRequestPreprocessor()
+     * \since QGIS 3.22
+     */
+#ifndef SIP_RUN
+    static QString setRequestPreprocessor( const std::function< void( QNetworkRequest *request )> &processor );
+#else
+    static QString setRequestPreprocessor( SIP_PYCALLABLE / AllowNone / );
+    % MethodCode
+    PyObject *s = 0;
+    Py_BEGIN_ALLOW_THREADS
+    Py_XINCREF( a0 );
+    QString id = QgsNetworkAccessManager::setRequestPreprocessor( [a0]( QNetworkRequest *arg )->QString
+    {
+      QString res;
+      SIP_BLOCK_THREADS
+      PyObject *s = sipCallMethod( NULL, a0, "D", arg, sipType_QNetworkRequest, NULL );
+      int state;
+      int sipIsError = 0;
+      QString *t1 = reinterpret_cast<QString *>( sipConvertToType( s, sipType_QString, 0, SIP_NOT_NONE, &state, &sipIsError ) );
+      if ( sipIsError == 0 )
+      {
+        res = QString( *t1 );
+      }
+      sipReleaseType( t1, sipType_QString, state );
+      SIP_UNBLOCK_THREADS
+      return res;
+    } );
+
+    s = sipConvertFromNewType( new QString( id ), sipType_QString, 0 );
+    Py_END_ALLOW_THREADS
+    return s;
+    % End
+#endif
+
+    /**
+     * Removes the custom pre-processor function with matching \a id.
+     *
+     * The \a id must correspond to a pre-processor previously added via a call to setRequestPreprocessor().
+     *
+     * Returns TRUE if processor existed and was removed.
+     *
+     * \see setRequestPreprocessor()
+     * \since QGIS 3.22
+     */
+#ifndef SIP_RUN
+    static bool removeRequestPreprocessor( const QString &id );
+#else
+    static void removeRequestPreprocessor( const QString &id );
+    % MethodCode
+    if ( !QgsNetworkAccessManager::removeRequestPreprocessor( *a0 ) )
+    {
+      PyErr_SetString( PyExc_KeyError, QStringLiteral( "No processor with id %1 exists." ).arg( *a0 ).toUtf8().constData() );
+      sipIsErr = 1;
+    }
+    % End
+#endif
+
     /**
      * Forwards an external browser login \a url opening request to the authentication handler.
      *
@@ -546,6 +611,12 @@ class CORE_EXPORT QgsNetworkAccessManager : public QNetworkAccessManager
     static const inline QgsSettingsEntryInteger settingsNetworkTimeout = QgsSettingsEntryInteger( QStringLiteral( "/qgis/networkAndProxy/networkTimeout" ), QgsSettings::NoSection, 60000, QObject::tr( "Network timeout" ) );
 #endif
 
+    /**
+     * Preprocesses request
+     * \param req the request to preprocess
+     * \since QGIS 3.22
+     */
+    void preprocessRequest( QNetworkRequest *req ) const;
 
   signals:
 
@@ -757,7 +828,6 @@ class CORE_EXPORT QgsNetworkAccessManager : public QNetworkAccessManager
     QMutex mAuthRequestHandlerMutex;
     // only in use by worker threads, unused in main thread
     QWaitCondition mAuthRequestWaitCondition;
-
 };
 
 #endif // QGSNETWORKACCESSMANAGER_H
diff --git a/src/core/qgstilecache.cpp b/src/core/qgstilecache.cpp
@@ -33,16 +33,21 @@ void QgsTileCache::insertTile( const QUrl &url, const QImage &image )
 
 bool QgsTileCache::tile( const QUrl &url, QImage &image )
 {
+  QNetworkRequest req( url );
+  //Preprocessing might alter the url, so we need to make sure we store/retrieve the url after preprocessing
+  QgsNetworkAccessManager::instance()->preprocessRequest( &req );
+  QUrl adjUrl = req.url();
+
   QMutexLocker locker( &sTileCacheMutex );
   bool success = false;
-  if ( QImage *i = sTileCache.object( url ) )
+  if ( QImage *i = sTileCache.object( adjUrl ) )
   {
     image = *i;
     success = true;
   }
-  else if ( QgsNetworkAccessManager::instance()->cache()->metaData( url ).isValid() )
+  else if ( QgsNetworkAccessManager::instance()->cache()->metaData( adjUrl ).isValid() )
   {
-    if ( QIODevice *data = QgsNetworkAccessManager::instance()->cache()->data( url ) )
+    if ( QIODevice *data = QgsNetworkAccessManager::instance()->cache()->data( adjUrl ) )
     {
       QByteArray imageData = data->readAll();
       delete data;
@@ -53,7 +58,7 @@ bool QgsTileCache::tile( const QUrl &url, QImage &image )
       // Check for null because it could be a redirect (see: https://github.com/qgis/QGIS/issues/24336 )
       if ( ! image.isNull( ) )
       {
-        sTileCache.insert( url, new QImage( image ) );
+        sTileCache.insert( adjUrl, new QImage( image ) );
         success = true;
       }
     }

diff --git a/tests/src/core/testqgsnetworkaccessmanager.cpp b/tests/src/core/testqgsnetworkaccessmanager.cpp
@@ -150,6 +150,7 @@ class TestQgsNetworkAccessManager : public QObject
     void cleanupTestCase();// will be called after the last testfunction was executed.
     void init();// will be called before each testfunction is executed.
     void cleanup();// will be called after every testfunction.
+    void testRequestPreprocessor();
     void testProxyExcludeList();
     void fetchEmptyUrl(); //test fetching blank url
     void fetchBadUrl(); //test fetching bad url
@@ -1127,8 +1128,17 @@ void TestQgsNetworkAccessManager::testCookieManagement()
   thread2.start();
   evLoop.exec();
   QVERIFY( thread2.getResult() );
+};
 
-}
+void TestQgsNetworkAccessManager::testRequestPreprocessor()
+{
+  QString processorId = QgsNetworkAccessManager::instance()->setRequestPreprocessor( []( QNetworkRequest * request ) { request->setHeader( QNetworkRequest::UserAgentHeader, QStringLiteral( "QGIS" ) );} );
+  QNetworkRequest request;
+  QgsNetworkAccessManager::instance()->preprocessRequest( &request );
+  QString userAgent = request.header( QNetworkRequest::UserAgentHeader ).toString();
+  QCOMPARE( userAgent, "QGIS" );
+  QgsNetworkAccessManager::instance()->removeRequestPreprocessor( processorId );
+};
 
 QGSTEST_MAIN( TestQgsNetworkAccessManager )
 #include "testqgsnetworkaccessmanager.moc"