A new way of hooking internet protocols.

chrome_frame/protocol_sink_wrap.h

Index: chrome_frame/protocol_sink_wrap.h
===================================================================
--- chrome_frame/protocol_sink_wrap.h	(revision 48838)
+++ chrome_frame/protocol_sink_wrap.h	(working copy)
@@ -16,6 +16,7 @@
 #include "base/basictypes.h"
 #include "base/ref_counted.h"
 #include "base/scoped_comptr_win.h"
+#include "base/scoped_bstr_win.h"
 #include "googleurl/src/gurl.h"
 #include "chrome_frame/chrome_frame_delegate.h"
 #include "chrome_frame/ie8_types.h"
@@ -37,6 +38,15 @@
 typedef HRESULT (STDMETHODCALLTYPE* InternetProtocolRoot_Continue_Fn)(
     IInternetProtocolRoot* me, PROTOCOLDATA* data);
 
+
+enum RendererType {
+  UNDETERMINED,
+  CHROME,
+  OTHER
+};
+
+class ProtData;
+
 // A class to wrap protocol sink in IInternetProtocol::Start[Ex] for
 // HTTP and HTTPS protocols.
 //
@@ -50,199 +60,101 @@
 // but delegate simply the QI.
 class ProtocolSinkWrap
     : public CComObjectRootEx<CComMultiThreadModel>,
-      public IInternetProtocolSink,
-      public IInternetBindInfoEx,
-      public IServiceProvider,
-      public IAuthenticate,
-      public IInternetProtocolEx,
-      public IInternetPriority,
-      public IWrappedProtocol,
-      // public IPreBindingSupport, // undocumented
-      // public ITransProtocolSink, // Undocumented
-      // public ITransactionInternal, // undocumented
-      public IUriContainer {
+      public IInternetProtocolSink {
  public:
 
 BEGIN_COM_MAP(ProtocolSinkWrap)
   COM_INTERFACE_ENTRY(IInternetProtocolSink)
-  COM_INTERFACE_ENTRY(IInternetBindInfo)
-  COM_INTERFACE_ENTRY(IInternetBindInfoEx)
-  COM_INTERFACE_ENTRY(IServiceProvider)
-  COM_INTERFACE_ENTRY(IAuthenticate)
-  COM_INTERFACE_ENTRY(IInternetProtocolRoot)
-  COM_INTERFACE_ENTRY(IInternetProtocol)
-  COM_INTERFACE_ENTRY(IInternetProtocolEx)
-  COM_INTERFACE_ENTRY(IInternetPriority)
-  COM_INTERFACE_ENTRY(IWrappedProtocol)
-  COM_INTERFACE_ENTRY_IF_DELEGATE_SUPPORTS(IUriContainer)
   COM_INTERFACE_BLIND_DELEGATE()
 END_COM_MAP()
 
-  ProtocolSinkWrap();
-  virtual ~ProtocolSinkWrap();
+  static ScopedComPtr<IInternetProtocolSink> CreateNewSink(
+      IInternetProtocolSink* sink, ProtData* prot_data);
 
-  bool Initialize(IInternetProtocol* protocol,
-      IInternetProtocolSink* original_sink, const wchar_t* url);
+  // Apparently this has to be public, to satisfy COM_INTERFACE_BLIND_DELEGATE
+  IInternetProtocolSink* delegate() {
+    return delegate_;
+  }
 
-  // VTable patches the IInternetProtocol and IIntenetProtocolEx interface.
-  // Returns true on success.
-  static bool PatchProtocolHandlers();
+ protected:
+  ProtocolSinkWrap();
+  ~ProtocolSinkWrap();
 
-  // Unpatches the IInternetProtocol and IInternetProtocolEx interfaces.
-  static void UnpatchProtocolHandlers();
-
-  // IInternetProtocol/Ex patches.
-  static STDMETHODIMP OnStart(InternetProtocol_Start_Fn orig_start,
-      IInternetProtocol* protocol, LPCWSTR url,
-      IInternetProtocolSink* prot_sink, IInternetBindInfo* bind_info,
-      DWORD flags, HANDLE_PTR reserved);
-
-  static STDMETHODIMP OnStartEx(
-      InternetProtocol_StartEx_Fn orig_start_ex, IInternetProtocolEx* protocol,
-      IUri* uri, IInternetProtocolSink* prot_sink,
-      IInternetBindInfo* bind_info, DWORD flags, HANDLE_PTR reserved);
-
-  static STDMETHODIMP OnRead(InternetProtocol_Read_Fn orig_read,
-      IInternetProtocol* protocol, void* buffer, ULONG size, ULONG* size_read);
-
+ private:
   // IInternetProtocolSink methods
   STDMETHOD(Switch)(PROTOCOLDATA* protocol_data);
   STDMETHOD(ReportProgress)(ULONG status_code, LPCWSTR status_text);
   STDMETHOD(ReportData)(DWORD flags, ULONG progress, ULONG max_progress);
   STDMETHOD(ReportResult)(HRESULT result, DWORD error, LPCWSTR result_text);
 
-  // IInternetBindInfoEx
-  STDMETHOD(GetBindInfo)(DWORD* flags, BINDINFO* bind_info);
-  STDMETHOD(GetBindString)(ULONG string_type, LPOLESTR* string_array,
-      ULONG array_size, ULONG* size_returned);
-  STDMETHOD(GetBindInfoEx)(DWORD *flags, BINDINFO* bind_info,
-      DWORD* bindf2, DWORD *reserved);
+  // Remember original sink
+  ScopedComPtr<IInternetProtocolSink> delegate_;
+  scoped_refptr<ProtData> prot_data_;
+  DISALLOW_COPY_AND_ASSIGN(ProtocolSinkWrap);
+};
 
-  // IServiceProvider
-  STDMETHOD(QueryService)(REFGUID service_guid, REFIID riid, void** service);
+class ProtData : public base::RefCounted<ProtData> {
+ public:
+  ProtData(IInternetProtocol* protocol, InternetProtocol_Read_Fn read_fun,
+           const wchar_t* url);
+  ~ProtData();
+  HRESULT Read(void* buffer, ULONG size, ULONG* size_read);
+  HRESULT ReportProgress(IInternetProtocolSink* delegate,
+                         ULONG status_code,
+                         LPCWSTR status_text);
+  HRESULT ReportData(IInternetProtocolSink* delegate,
+                     DWORD flags, ULONG progress, ULONG max_progress);
+  HRESULT ReportResult(IInternetProtocolSink* delegate, HRESULT result,
+                       DWORD error, LPCWSTR result_text);
+  void UpdateUrl(const wchar_t* url);
+  static scoped_refptr<ProtData> DataFromProtocol(IInternetProtocol* protocol);
 
-  // IAuthenticate
-  STDMETHOD(Authenticate)(HWND* window, LPWSTR* user_name, LPWSTR* password);
-
-  // IInternetProtocolEx
-  STDMETHOD(Start)(LPCWSTR url, IInternetProtocolSink *protocol_sink,
-      IInternetBindInfo* bind_info, DWORD flags, HANDLE_PTR reserved);
-  STDMETHOD(Continue)(PROTOCOLDATA* protocol_data);
-  STDMETHOD(Abort)(HRESULT reason, DWORD options);
-  STDMETHOD(Terminate)(DWORD options);
-  STDMETHOD(Suspend)();
-  STDMETHOD(Resume)();
-  STDMETHOD(Read)(void *buffer, ULONG size, ULONG* size_read);
-  STDMETHOD(Seek)(LARGE_INTEGER move, DWORD origin, ULARGE_INTEGER* new_pos);
-  STDMETHOD(LockRequest)(DWORD options);
-  STDMETHOD(UnlockRequest)();
-  STDMETHOD(StartEx)(IUri* uri, IInternetProtocolSink* protocol_sink,
-      IInternetBindInfo* bind_info, DWORD flags, HANDLE_PTR reserved);
-
-  // IInternetPriority
-  STDMETHOD(SetPriority)(LONG priority);
-  STDMETHOD(GetPriority)(LONG* priority);
-
-  // IWrappedProtocol
-  STDMETHOD(GetWrapperCode)(LONG *code, DWORD_PTR reserved);
-
-  // public IUriContainer
-  STDMETHOD(GetIUri)(IUri** uri);
-
-  // IPreBindingSupport, // undocumented
-  // ITransProtocolSink, // Undocumented
-  // ITransactionInternal, // undocumented
-
-  IInternetProtocolSink* delegate() const {
-    return delegate_;
-  }
-
- protected:
-  enum RendererType {
-    UNDETERMINED,
-    CHROME,
-    OTHER
-  };
-
-  typedef std::map<IInternetProtocol*, ProtocolSinkWrap*> ProtocolSinkMap;
-  static const int kMaxContentSniffLength = 1024;
-
-  static scoped_refptr<ProtocolSinkWrap> InstanceFromProtocol(
-      IInternetProtocol* protocol);
-  static ScopedComPtr<IInternetProtocolSink> MaybeWrapSink(
-      IInternetProtocol* protocol, IInternetProtocolSink* prot_sink,
-      const wchar_t* url);
-
-  void DetermineRendererType();
-
-  HRESULT OnReadImpl(void* buffer, ULONG size, ULONG* size_read,
-      InternetProtocol_Read_Fn orig_read);
-
-  // This function determines whether the current request should be handled
-  // by Chrome. If yes it reports the mime type as application/chromepage,
-  // which ensures that the ChromeFrame is instantiated for handling this
-  // request.
-  // Returns S_OK on success.
-  HRESULT CheckAndReportChromeMimeTypeForRequest();
-
-  bool is_undetermined() const {
-    return (UNDETERMINED == renderer_type_);
-  }
-  RendererType renderer_type() const {
+  RendererType renderer_type() {
     return renderer_type_;
   }
 
-  // Creates an instance of the specified protocol handler and returns the
-  // IInternetProtocol interface pointer.
-  // Returns S_OK on success.
-  static HRESULT CreateProtocolHandlerInstance(const CLSID& clsid,
-                                               IInternetProtocol** protocol);
+ private:
+  typedef std::map<IInternetProtocol*, ProtData*> ProtocolDataMap;
+  static ProtocolDataMap datamap_;
+  static Lock datamap_lock_;
 
-  // Helper function for patching the VTable of the IInternetProtocol
-  // interface. It instantiates the object identified by the protocol_clsid
-  // parameter and patches its VTable.
-  // Returns S_OK on success.
-  static HRESULT PatchProtocolMethods(
-      const CLSID& protocol_clsid,
-      vtable_patch::MethodPatchInfo* protocol_patch_info,
-      vtable_patch::MethodPatchInfo* protocol_ex_patch_info);
+  // Url we are retrieving. Used for IsOptInUrl() only.
+  std::wstring url_;
+  // Our gate to IInternetProtocol::Read()
+  IInternetProtocol* protocol_;
+  InternetProtocol_Read_Fn read_fun_;
 
-  // WARNING: Don't use GURL variables here.  Please see
-  // http://b/issue?id=2102171 for details.
+  // What BINDSTATUS_MIMETYPEAVAILABLE and Co. tells us.
+  ScopedBstr suggested_mime_type_;
+  // At least one of the following has been received:
+  // BINDSTATUS_MIMETYPEAVAILABLE,
+  // MIMESTATUS_VERIFIEDMIMETYPEAVAILABLE
+  // BINDSTATUS_SERVER_MIMETYPEAVAILABLE
+  bool has_suggested_mime_type_;
+  // BINDSTATUS_SERVER_MIMETYPEAVAILABLE received, so we shall fire one.
+  bool has_server_mime_type_;
 
-  // Remember original sink
-  ScopedComPtr<IInternetProtocolSink> delegate_;
-
-  // Cannot take a reference on the protocol.
-  IInternetProtocol* protocol_;
+  // Did we received ReportData()
+  bool report_data_received_;
   RendererType renderer_type_;
 
   // Buffer for accumulated data including 1 extra for NULL-terminator
+  static const size_t kMaxContentSniffLength = 2 * 1024;
   char buffer_[kMaxContentSniffLength + 1];
   unsigned long buffer_size_;  // NOLINT
   unsigned long buffer_pos_;  // NOLINT
 
-  // Accumulated result
-  bool is_saved_result_;
-  HRESULT result_code_;
-  DWORD result_error_;
-  std::wstring result_text_;
-  // For tracking re-entrency and preventing duplicate Read()s from
-  // distorting the outcome of ReportData.
-  int report_data_recursiveness_;
+  HRESULT FillBuffer();
+  void SaveSuggestedMimeType(LPCWSTR status_text);
+  void FireSugestedMimeType(IInternetProtocolSink* delegate);
+};
 
-  static ProtocolSinkMap sink_map_;
-  // TODO(joshia): Replace with Lock
-  static CComAutoCriticalSection sink_map_lock_;
-
-  std::wstring url_;
-
-  // Set to true if we are in the context of determining the desired renderer
-  // type.
-  bool determining_renderer_type_;
-
- private:
-  DISALLOW_COPY_AND_ASSIGN(ProtocolSinkWrap);
+struct TransactionHooks {
+  void InstallHooks();
+  void RevertHooks();
 };
 
+DECLSPEC_SELECTANY struct TransactionHooks g_trans_hooks;
+
 #endif  // CHROME_FRAME_PROTOCOL_SINK_WRAP_H_
+